integrated naga transpiling to wgsl

Author: Firestar99

Cargo.lock

@@ -28,6 +28,7 @@ tempdir = "0.3.7"
 test-log = "0.2.16"
 cargo_metadata = "0.19.2"
 semver = "1.0.26"
+naga = "25.0.1"
 
 # This crate MUST NEVER be upgraded, we need this particular "first" version to support old rust-gpu builds
 legacy_target_specs = { package = "rustc_codegen_spirv-target-specs", version = "0.9.0", features = ["include_str"] }

Cargo.toml

@@ -10,6 +10,10 @@ build = "build.rs"
 default-run = "cargo-gpu"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
+[features]
+# Enable naga support to convert spirv to wgsl
+naga = ["dep:naga"]
+
 [dependencies]
 cargo_metadata.workspace = true
 anyhow.workspace = true
@@ -24,6 +28,7 @@ serde.workspace = true
 serde_json.workspace = true
 crossterm.workspace = true
 semver.workspace = true
+naga = { workspace = true, optional = true, features = ["spv-in", "wgsl-out"] }
 
 [dev-dependencies]
 test-log.workspace = true

crates/cargo-gpu/Cargo.toml

@@ -65,11 +65,15 @@ mod legacy_target_specs;
 mod linkage;
 mod lockfile;
 mod metadata;
+#[cfg(feature = "naga")]
+mod naga;
 mod show;
 mod spirv_source;
 mod test;
 
 pub use install::*;
+#[cfg(feature = "naga")]
+pub use naga::*;
 pub use spirv_builder;
 
 /// Central function to write to the user.

crates/cargo-gpu/src/lib.rs

@@ -23,8 +23,13 @@ impl Linkage {
                 .map(|comp| comp.as_os_str().to_string_lossy())
                 .collect::<Vec<_>>()
                 .join("/"),
-            wgsl_entry_point: entry_point.as_ref().replace("::", ""),
+            wgsl_entry_point: spv_entry_point_to_wgsl(entry_point.as_ref()),
             entry_point: entry_point.as_ref().to_owned(),
         }
     }
 }
+
+/// Convert a spirv entry point to a valid wgsl entry point
+pub fn spv_entry_point_to_wgsl(entry_point: &str) -> String {
+    entry_point.replace("::", "")
+}

crates/cargo-gpu/src/linkage.rs

@@ -0,0 +1,98 @@
+//! naga transpiling to wgsl support, hidden behind feature `naga`
+
+use crate::linkage::spv_entry_point_to_wgsl;
+use anyhow::Context as _;
+use naga::error::ShaderError;
+pub use naga::valid::Capabilities;
+use spirv_builder::{CompileResult, ModuleResult};
+use std::collections::BTreeMap;
+use std::path::{Path, PathBuf};
+
+/// convert a single spv file to a wgsl file using naga
+fn spv_to_wgsl(spv_src: &Path, wgsl_dst: &Path, capabilities: Capabilities) -> anyhow::Result<()> {
+    let inner = || -> anyhow::Result<()> {
+        let spv_bytes = std::fs::read(spv_src).context("could not read spv file")?;
+        let opts = naga::front::spv::Options::default();
+        let module = naga::front::spv::parse_u8_slice(&spv_bytes, &opts)
+            .map_err(|err| ShaderError {
+                source: String::new(),
+                label: None,
+                inner: Box::new(err),
+            })
+            .context("naga could not parse spv")?;
+        let mut validator =
+            naga::valid::Validator::new(naga::valid::ValidationFlags::default(), capabilities);
+        let info = validator
+            .validate(&module)
+            .map_err(|err| ShaderError {
+                source: String::new(),
+                label: None,
+                inner: Box::new(err),
+            })
+            .context("validation of naga module failed")?;
+        let wgsl =
+            naga::back::wgsl::write_string(&module, &info, naga::back::wgsl::WriterFlags::empty())
+                .context("naga conversion to wgsl failed")?;
+        std::fs::write(wgsl_dst, wgsl).context("failed to write wgsl file")?;
+        Ok(())
+    };
+    inner().with_context(|| {
+        format!(
+            "converting spv '{}' to wgsl '{}' failed    ",
+            spv_src.display(),
+            wgsl_dst.display()
+        )
+    })
+}
+
+/// convert spv file path to a valid unique wgsl file path
+fn wgsl_file_name(path: &Path) -> PathBuf {
+    path.with_extension("wgsl")
+}
+
+/// Extension trait for naga transpiling
+pub trait CompileResultNagaExt {
+    /// Transpile the spirv binaries to wgsl source code using [`naga`], typically for webgpu compatibility.
+    ///
+    /// Converts this [`CompileResult`] of spirv binaries and entry points to a [`CompileResult`] pointing to wgsl source code files and their associated wgsl entry
+    /// points.
+    ///
+    /// # Errors
+    /// [`naga`] transpile may error in various ways
+    fn transpile_to_wgsl(&self, capabilities: Capabilities) -> anyhow::Result<CompileResult>;
+}
+
+impl CompileResultNagaExt for CompileResult {
+    #[inline]
+    fn transpile_to_wgsl(&self, capabilities: Capabilities) -> anyhow::Result<CompileResult> {
+        Ok(match &self.module {
+            ModuleResult::SingleModule(spv) => {
+                let wgsl = wgsl_file_name(spv);
+                spv_to_wgsl(spv, &wgsl, capabilities)?;
+                let entry_points = self
+                    .entry_points
+                    .iter()
+                    .map(|entry| spv_entry_point_to_wgsl(entry))
+                    .collect();
+                Self {
+                    entry_points,
+                    module: ModuleResult::SingleModule(wgsl),
+                }
+            }
+            ModuleResult::MultiModule(map) => {
+                let new_map: BTreeMap<String, PathBuf> = map
+                    .iter()
+                    .map(|(entry_point, spv)| {
+                        let wgsl = wgsl_file_name(spv);
+                        spv_to_wgsl(spv, &wgsl, capabilities)?;
+                        Ok((spv_entry_point_to_wgsl(entry_point), wgsl))
+                    })
+                    .collect::<anyhow::Result<_>>()?;
+                Self {
+                    entry_points: new_map.keys().cloned().collect(),
+                    module: ModuleResult::MultiModule(new_map),
+                }
+            }
+        })
+    }
+}