Build the wgpu shader at runtime (#627)

* Build the wgpu shader at runtime

.github/workflows/ci.yaml

@@ -23,6 +23,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     env:
       spirv_tools_version: "20200928"
+      RUSTUP_UNPACK_RAM: "104857600"
     steps:
       - uses: actions/checkout@v2
       # Ubuntu does have `brew install spirv-tools`, but it installs from

Cargo.lock

@@ -783,6 +783,13 @@ dependencies = [
  "winit",
 ]
 
+[[package]]
+name = "example-runner-wgpu-builder"
+version = "0.4.0-alpha.7"
+dependencies = [
+ "spirv-builder",
+]
+
 [[package]]
 name = "filetime"
 version = "0.2.14"

Cargo.toml

@@ -3,6 +3,7 @@ members = [
     "examples/runners/cpu",
     "examples/runners/ash",
     "examples/runners/wgpu",
+    "examples/runners/wgpu/builder",
     "examples/shaders/sky-shader",
     "examples/shaders/simplest-shader",
     "examples/shaders/compute-shader",

crates/rustc_codegen_spirv/build.rs

@@ -1,8 +0,0 @@
-// Putting this check here causes compilation failure seconds into the build,
-// putting it in lib.rs fails after minutes because spirv-tools gets compiled first.
-#[cfg(all(feature = "use-compiled-tools", feature = "use-installed-tools"))]
-compile_error!(
-    "Either \"use-compiled-tools\" (enabled by default) or \"use-installed-tools\" may be enabled."
-);
-
-fn main() {}

crates/rustc_codegen_spirv/src/lib.rs

@@ -80,6 +80,17 @@
 )]
 #![deny(clippy::unimplemented, clippy::ok_expect)]
 
+// Unfortunately, this will not fail fast when compiling, but rather will wait for
+// rustc_codegen_spirv to be compiled. Putting this in build.rs will solve that problem, however,
+// that creates the much worse problem that then running `cargo check` will cause
+// rustc_codegen_spirv to be *compiled* instead of merely checked, something that takes
+// significantly longer. So, the trade-off between detecting a bad configuration slower for a
+// faster `cargo check` is worth it.
+#[cfg(all(feature = "use-compiled-tools", feature = "use-installed-tools"))]
+compile_error!(
+    "Either \"use-compiled-tools\" (enabled by default) or \"use-installed-tools\" may be enabled."
+);
+
 extern crate rustc_ast;
 extern crate rustc_attr;
 extern crate rustc_codegen_ssa;

examples/runners/wgpu/Cargo.toml

@@ -24,7 +24,7 @@ winit = { version = "0.24", features = ["web-sys"] }
 clap = "3.0.0-beta.2"
 strum = { version = "0.20", default_features = false, features = ["derive"] }
 
-[build-dependencies]
+[target.'cfg(not(any(target_os = "android", target_arch = "wasm32")))'.dependencies]
 spirv-builder = { path = "../../../crates/spirv-builder", default-features = false }
 
 [target.'cfg(target_os = "android")'.dependencies]

examples/runners/wgpu/build.rs

@@ -1,15 +1,52 @@
-use spirv_builder::SpirvBuilder;
+use std::env;
 use std::error::Error;
-
-fn build_shader(path_to_create: &str) -> Result<(), Box<dyn Error>> {
-    SpirvBuilder::new(path_to_create, "spirv-unknown-vulkan1.0").build()?;
-    Ok(())
-}
+use std::path::PathBuf;
 
 fn main() -> Result<(), Box<dyn Error>> {
-    build_shader("../../shaders/sky-shader")?;
-    build_shader("../../shaders/simplest-shader")?;
-    build_shader("../../shaders/compute-shader")?;
-    build_shader("../../shaders/mouse-shader")?;
+    let target_os = std::env::var("CARGO_CFG_TARGET_OS")?;
+    let target_arch = std::env::var("CARGO_CFG_TARGET_ARCH")?;
+    // Always build on CI to make sure the shaders can still build
+    let is_on_ci = std::env::var("CI");
+    println!("cargo:rerun-if-changed=build.rs");
+    println!("cargo:rerun-if-env-changed=CARGO_CFG_TARGET_OS");
+    println!("cargo:rerun-if-env-changed=CARGO_CFG_TARGET_ARCH");
+    println!("cargo:rerun-if-env-changed=CI");
+    // While OUT_DIR is set for both build.rs and compiling the crate, PROFILE is only set in
+    // build.rs. So, export it to crate compilation as well.
+    let profile = env::var("PROFILE").unwrap();
+    println!("cargo:rustc-env=PROFILE={}", profile);
+    if target_os != "android" && target_arch != "wasm32" && is_on_ci.is_err() {
+        return Ok(());
+    }
+    let mut dir = PathBuf::from(env::var_os("OUT_DIR").unwrap());
+    // Strip `$profile/build/*/out`.
+    let ok = dir.ends_with("out")
+        && dir.pop()
+        && dir.pop()
+        && dir.ends_with("build")
+        && dir.pop()
+        && dir.ends_with(profile)
+        && dir.pop();
+    assert!(ok);
+    let dir = dir.join("spirv-builder");
+    let status = std::process::Command::new("cargo")
+        .args([
+            "run",
+            "--release",
+            "-p",
+            "example-runner-wgpu-builder",
+            "--target-dir",
+        ])
+        .arg(dir)
+        .stderr(std::process::Stdio::inherit())
+        .stdout(std::process::Stdio::inherit())
+        .status()?;
+    if !status.success() {
+        if let Some(code) = status.code() {
+            std::process::exit(code);
+        } else {
+            std::process::exit(1);
+        }
+    }
     Ok(())
 }

examples/runners/wgpu/builder/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "example-runner-wgpu-builder"
+version = "0.4.0-alpha.7"
+authors = ["Embark <opensource@embark-studios.com>"]
+edition = "2018"
+license = "MIT OR Apache-2.0"
+publish = false
+
+# See rustc_codegen_spirv/Cargo.toml for details on these features
+[features]
+default = ["use-compiled-tools"]
+use-installed-tools = ["spirv-builder/use-installed-tools"]
+use-compiled-tools = ["spirv-builder/use-compiled-tools"]
+
+[dependencies]
+spirv-builder = { path = "../../../../crates/spirv-builder", default-features = false }

examples/runners/wgpu/builder/src/main.rs

@@ -0,0 +1,15 @@
+use spirv_builder::SpirvBuilder;
+use std::error::Error;
+
+fn build_shader(path_to_create: &str) -> Result<(), Box<dyn Error>> {
+    SpirvBuilder::new(path_to_create, "spirv-unknown-vulkan1.0").build()?;
+    Ok(())
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
+    build_shader("../../shaders/sky-shader")?;
+    build_shader("../../shaders/simplest-shader")?;
+    build_shader("../../shaders/compute-shader")?;
+    build_shader("../../shaders/mouse-shader")?;
+    Ok(())
+}

examples/runners/wgpu/src/compute.rs

@@ -34,10 +34,12 @@ fn create_device_queue() -> (wgpu::Device, wgpu::Queue) {
 }
 
 pub fn start(options: &Options) {
+    let shader_binary = shader_module(options.shader);
+
     let (device, queue) = create_device_queue();
 
     // Load the shaders from disk
-    let module = device.create_shader_module(&shader_module(options.shader));
+    let module = device.create_shader_module(&shader_binary);
 
     let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
         label: None,

examples/runners/wgpu/src/graphics.rs

@@ -20,10 +20,10 @@ fn mouse_button_index(button: MouseButton) -> usize {
 }
 
 async fn run(
-    options: &Options,
     event_loop: EventLoop<()>,
     window: Window,
     swapchain_format: wgpu::TextureFormat,
+    shader_binary: wgpu::ShaderModuleDescriptor<'_>,
 ) {
     let size = window.inner_size();
     let instance = wgpu::Instance::new(wgpu::BackendBit::VULKAN | wgpu::BackendBit::METAL);
@@ -65,7 +65,7 @@ async fn run(
         .expect("Failed to create device");
 
     // Load the shaders from disk
-    let module = device.create_shader_module(&shader_module(options.shader));
+    let module = device.create_shader_module(&shader_binary);
 
     let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
         label: None,
@@ -273,6 +273,9 @@ async fn run(
 }
 
 pub fn start(options: &Options) {
+    // Build the shader before we pop open a window, since it might take a while.
+    let shader_binary = shader_module(options.shader);
+
     let event_loop = EventLoop::new();
     let window = winit::window::WindowBuilder::new()
         .with_title("Rust GPU - wgpu")
@@ -303,7 +306,6 @@ pub fn start(options: &Options) {
         } else {
             wgpu_subscriber::initialize_default_subscriber(None);
             futures::executor::block_on(run(
-                options,
                 event_loop,
                 window,
                 if cfg!(target_os = "android") {
@@ -311,6 +313,7 @@ pub fn start(options: &Options) {
                 } else {
                     wgpu::TextureFormat::Bgra8UnormSrgb
                 },
+                shader_binary,
             ));
         }
     }

examples/runners/wgpu/src/lib.rs

@@ -57,6 +57,55 @@ pub enum RustGPUShader {
 }
 
 fn shader_module(shader: RustGPUShader) -> wgpu::ShaderModuleDescriptor<'static> {
+    #[cfg(not(any(target_os = "android", target_arch = "wasm32")))]
+    {
+        use spirv_builder::SpirvBuilder;
+        use std::borrow::Cow;
+        use std::path::{Path, PathBuf};
+        // Hack: spirv_builder builds into a custom directory if running under cargo, to not
+        // deadlock, and the default target directory if not. However, packages like `proc-macro2`
+        // have different configurations when being built here vs. when building
+        // rustc_codegen_spirv normally, so we *want* to build into a separate target directory, to
+        // not have to rebuild half the crate graph every time we run. So, pretend we're running
+        // under cargo by setting these environment variables.
+        std::env::set_var("OUT_DIR", env!("OUT_DIR"));
+        std::env::set_var("PROFILE", env!("PROFILE"));
+        let crate_name = match shader {
+            RustGPUShader::Simplest => "sky-shader",
+            RustGPUShader::Sky => "simplest-shader",
+            RustGPUShader::Compute => "compute-shader",
+            RustGPUShader::Mouse => "mouse-shader",
+        };
+        let manifest_dir = env!("CARGO_MANIFEST_DIR");
+        let crate_path = [
+            Path::new(manifest_dir),
+            Path::new(".."),
+            Path::new(".."),
+            Path::new("shaders"),
+            Path::new(crate_name),
+        ]
+        .iter()
+        .copied()
+        .collect::<PathBuf>();
+        let result = SpirvBuilder::new(crate_path, "spirv-unknown-vulkan1.0")
+            .print_metadata(false)
+            .build()
+            .unwrap();
+        let data = std::fs::read(result).unwrap();
+        let spirv = wgpu::util::make_spirv(&data);
+        let spirv = match spirv {
+            wgpu::ShaderSource::Wgsl(cow) => wgpu::ShaderSource::Wgsl(Cow::Owned(cow.into_owned())),
+            wgpu::ShaderSource::SpirV(cow) => {
+                wgpu::ShaderSource::SpirV(Cow::Owned(cow.into_owned()))
+            }
+        };
+        wgpu::ShaderModuleDescriptor {
+            label: None,
+            source: spirv,
+            flags: wgpu::ShaderFlags::default(),
+        }
+    }
+    #[cfg(any(target_os = "android", target_arch = "wasm32"))]
     match shader {
         RustGPUShader::Simplest => wgpu::include_spirv!(env!("simplest_shader.spv")),
         RustGPUShader::Sky => wgpu::include_spirv!(env!("sky_shader.spv")),