EmbarkStudios · eddyb · Dec 9, 2022 · Dec 1, 2022 · Dec 2, 2022 · Dec 2, 2022
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -151,6 +151,27 @@ jobs:
         run: RUSTDOCFLAGS=-Dwarnings cargo doc --no-deps
       - name: Clippy
         run: .github/workflows/clippy.sh
+      # HACK(eddyb) see `docs/src/codegen-args.md` for more context around this,
+      # but basically we're implementing a custom "lint" to ban `std::env` usage,
+      # which could be disastrous because env vars access can't be tracked by
+      # `rustc`, unlike its CLI flags (which are integrated with incremental).
+      - name: Disallow `std::env` (mis)use from `rustc_codegen_spirv``
+        run: |
+          if (
+              egrep -r '::\s*env|env\s*::' crates/rustc_codegen_spirv/src/ |
+              # HACK(eddyb) exclude the one place in `rustc_codegen_spirv`
+              # needing access to an env var (only for codegen args `--help`).
+              egrep -v '^crates/rustc_codegen_spirv/src/codegen_cx/mod.rs:            let help_flag_comes_from_spirv_builder_env_var = std::env::var\(spirv_builder_env_var\)$'
+          ); then
+            echo '^^^'
+            echo '!!! Found disallowed `std::env` usage in `rustc_codegen_spirv` !!!'
+            echo ' ("codegen args" should be used instead of environment variables)'
+            echo
+            echo 'For more details on "codegen args", see: docs/src/codegen-args.md'
+            echo ' (and/or https://github.com/EmbarkStudios/rust-gpu/pull/959)'
+            echo
+            exit 1
+          fi
 
   cargo-deny:
     runs-on: ubuntu-20.04

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -29,10 +29,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added ⭐
+
+- Added two `spirv-builder` environment variables to customize *only* the `rustc` invocations for shader crates and their dependencies:
+    - `RUSTGPU_RUSTFLAGS="..."` for shader `RUSTFLAGS="..."`
+    - `RUSTGPU_CODEGEN_ARGS="..."` for shader "codegen args" (i.e. `RUSTFLAGS=-Cllvm-args="..."`)  
+      (check out ["codegen args" docs](docs/src/codegen-args.md), or run with `RUSTGPU_CODEGEN_ARGS=--help` to see the full list of options)
+
 ### Changed 🛠️
 
 - Updated toolchain to `nightly-2022-10-29`
-- Applied workspace inheritance to Cargo.toml files
+- Applied workspace inheritance to `Cargo.toml` files
+- Moved `rustc_codegen_spirv` debugging functionality from environment variables to "codegen args" options/flags (see [the updated docs](docs/src/codegen-args.md) for more details)
 
 ### Removed 🔥
 
@@ -51,7 +59,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added ⭐
 
-- Added check for env var `RUSTGPU_SKIP_TOOLCHAIN_CHECK` to prevent toolchain check
+- Added check for environment variable `RUSTGPU_SKIP_TOOLCHAIN_CHECK` to prevent toolchain check
 
 ### Changed 🛠️
 

diff --git a/crates/rustc_codegen_spirv/src/codegen_cx/mod.rs b/crates/rustc_codegen_spirv/src/codegen_cx/mod.rs
@@ -31,8 +31,9 @@ use rustc_target::abi::call::{FnAbi, PassMode};
 use rustc_target::abi::{HasDataLayout, TargetDataLayout};
 use rustc_target::spec::{HasTargetSpec, Target};
 use std::cell::{Cell, RefCell};
+use std::collections::BTreeSet;
 use std::iter::once;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::rc::Rc;
 use std::str::FromStr;
 
@@ -231,8 +232,9 @@ impl<'tcx> CodegenCx<'tcx> {
     }
 }
 
-#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+#[derive(Debug, Default, PartialEq, Eq, Clone, Copy)]
 pub enum SpirvMetadata {
+    #[default]
     None,
     NameVariables,
     Full,
@@ -247,6 +249,8 @@ pub struct CodegenArgs {
 
     pub spirv_metadata: SpirvMetadata,
 
+    pub run_spirv_val: bool,
+
     // spirv-val flags
     pub relax_struct_store: bool,
     pub relax_logical_pointer: bool,
@@ -255,8 +259,22 @@ pub struct CodegenArgs {
     pub scalar_block_layout: bool,
     pub skip_block_layout: bool,
 
+    pub run_spirv_opt: bool,
+
     // spirv-opt flags
     pub preserve_bindings: bool,
+
+    /// All options pertinent to `rustc_codegen_spirv::linker` specifically.
+    //
+    // FIXME(eddyb) should these be handled as `-C linker-args="..."` instead?
+    pub linker_opts: crate::linker::Options,
+
+    // NOTE(eddyb) these are debugging options that used to be env vars
+    // (for more information see `docs/src/codegen-args.md`).
+    pub dump_mir: Option<PathBuf>,
+    pub dump_module_on_panic: Option<PathBuf>,
+    pub dump_pre_link: Option<PathBuf>,
+    pub dump_post_link: Option<PathBuf>,
 }
 
 impl CodegenArgs {
@@ -267,42 +285,179 @@ impl CodegenArgs {
         }
     }
 
+    // FIXME(eddyb) `structopt` would come a long way to making this nicer.
     pub fn parse(args: &[String]) -> Result<Self, rustc_session::getopts::Fail> {
         use rustc_session::getopts;
+
+        // FIXME(eddyb) figure out what casing ("Foo bar" vs "foo bar") to use
+        // for the descriptions, `rustc` seems a bit inconsistent itself on this.
+
         let mut opts = getopts::Options::new();
+        opts.optflag("h", "help", "Display this message");
         opts.optopt(
             "",
             "module-output",
             "single output or multiple output",
             "[single|multiple]",
         );
-        opts.optflagopt("", "disassemble", "print module to stderr", "");
+        opts.optflag("", "disassemble", "print module to stderr");
         opts.optopt("", "disassemble-fn", "print function to stderr", "NAME");
         opts.optopt(
             "",
             "disassemble-entry",
             "print entry point to stderr",
             "NAME",
         );
-        opts.optflagopt("", "disassemble-globals", "print globals to stderr", "");
+        opts.optflag("", "disassemble-globals", "print globals to stderr");
 
         opts.optopt("", "spirv-metadata", "how much metadata to include", "");
 
-        opts.optflagopt("", "relax-struct-store", "Allow store from one struct type to a different type with compatible layout and members.", "");
-        opts.optflagopt("", "relax-logical-pointer", "Allow allocating an object of a pointer type and returning a pointer value from a function in logical addressing mode", "");
-        opts.optflagopt("", "relax-block-layout", "Enable VK_KHR_relaxed_block_layout when checking standard uniform, storage buffer, and push constant layouts. This is the default when targeting Vulkan 1.1 or later.", "");
-        opts.optflagopt("", "uniform-buffer-standard-layout", "Enable VK_KHR_uniform_buffer_standard_layout when checking standard uniform buffer layouts.", "");
-        opts.optflagopt("", "scalar-block-layout", "Enable VK_EXT_scalar_block_layout when checking standard uniform, storage buffer, and push constant layouts. Scalar layout rules are more permissive than relaxed block layout so in effect this will override the --relax-block-layout option.", "");
-        opts.optflagopt("", "skip-block-layout", "Skip checking standard uniform/storage buffer layout. Overrides any --relax-block-layout or --scalar-block-layout option.", "");
+        // FIXME(eddyb) clean up this `no-` "negation prefix" situation.
+        opts.optflag(
+            "",
+            "no-spirv-val",
+            "disables running spirv-val on the final output",
+        );
+
+        opts.optflag("", "relax-struct-store", "Allow store from one struct type to a different type with compatible layout and members.");
+        opts.optflag("", "relax-logical-pointer", "Allow allocating an object of a pointer type and returning a pointer value from a function in logical addressing mode");
+        opts.optflag("", "relax-block-layout", "Enable VK_KHR_relaxed_block_layout when checking standard uniform, storage buffer, and push constant layouts. This is the default when targeting Vulkan 1.1 or later.");
+        opts.optflag("", "uniform-buffer-standard-layout", "Enable VK_KHR_uniform_buffer_standard_layout when checking standard uniform buffer layouts.");
+        opts.optflag("", "scalar-block-layout", "Enable VK_EXT_scalar_block_layout when checking standard uniform, storage buffer, and push constant layouts. Scalar layout rules are more permissive than relaxed block layout so in effect this will override the --relax-block-layout option.");
+        opts.optflag("", "skip-block-layout", "Skip checking standard uniform/storage buffer layout. Overrides any --relax-block-layout or --scalar-block-layout option.");
+
+        // FIXME(eddyb) clean up this `no-` "negation prefix" situation.
+        opts.optflag(
+            "",
+            "no-spirv-opt",
+            "disables running spirv-opt on the final output",
+        );
 
-        opts.optflagopt(
+        opts.optflag(
             "",
             "preserve-bindings",
             "Preserve unused descriptor bindings. Useful for reflection.",
+        );
+
+        // Linker options.
+        // FIXME(eddyb) should these be handled as `-C linker-args="..."` instead?
+        {
+            // FIXME(eddyb) clean up this `no-` "negation prefix" situation.
+            opts.optflag("", "no-dce", "disables running dead code elimination");
+            opts.optflag(
+                "",
+                "no-compact-ids",
+                "disables compaction of SPIR-V IDs at the end of linking",
+            );
+            opts.optflag("", "no-structurize", "disables CFG structurization");
+
+            // NOTE(eddyb) these are debugging options that used to be env vars
+            // (for more information see `docs/src/codegen-args.md`).
+            opts.optopt(
+                "",
+                "dump-post-merge",
+                "dump the merged module immediately after merging, to FILE",
+                "FILE",
+            );
+            opts.optopt(
+                "",
+                "dump-post-split",
+                "dump modules immediately after multimodule splitting, to files in DIR",
+                "DIR",
+            );
+            opts.optflag(
+                "",
+                "specializer-debug",
+                "enable debug logging for the specializer",
+            );
+            opts.optopt(
+                "",
+                "specializer-dump-instances",
+                "dump all instances inferred by the specializer, to FILE",
+                "FILE",
+            );
+            opts.optflag("", "print-all-zombie", "prints all removed zombies");
+            opts.optflag(
+                "",
+                "print-zombie",
+                "prints everything removed (even transitively) due to zombies",
+            );
+        }
+
+        // NOTE(eddyb) these are debugging options that used to be env vars
+        // (for more information see `docs/src/codegen-args.md`).
+        opts.optopt(
+            "",
+            "dump-mir",
+            "dump every MIR body codegen sees, to files in DIR",
+            "DIR",
+        );
+        opts.optopt(
+            "",
+            "dump-module-on-panic",
+            "if codegen panics, dump the (partially) emitted module, to FILE",
+            "FILE",
+        );
+        opts.optopt(
+            "",
+            "dump-pre-link",
+            "dump all input modules to the linker, to files in DIR",
+            "DIR",
+        );
+        opts.optopt(
             "",
+            "dump-post-link",
+            "dump all output modules from the linker, to files in DIR",
+            "DIR",
         );
 
         let matches = opts.parse(args)?;
+
+        let help_flag_positions: BTreeSet<_> = ["h", "help"]
+            .iter()
+            .flat_map(|&name| matches.opt_positions(name))
+            .collect();
+        if !help_flag_positions.is_empty() {
+            // HACK(eddyb) this tries to be a bit nicer to end-users, when they
+            // use `spirv-builder` (and so the `RUSTGPU_CODEGEN_ARGS` env var,
+            // to set codegen args), as mentioning `-Cllvm-args` is suboptimal.
+            let spirv_builder_env_var = "RUSTGPU_CODEGEN_ARGS";
+            let help_flag_comes_from_spirv_builder_env_var = std::env::var(spirv_builder_env_var)
+                .ok()
+                .and_then(|args_from_env| {
+                    let args_from_env: Vec<_> = args_from_env.split_whitespace().collect();
+                    if args_from_env.is_empty() {
+                        return None;
+                    }
+
+                    // HACK(eddyb) this may be a bit inefficient but we want to
+                    // make sure that *at least one* of the `-h`/`--help` flags
+                    // came from the `spirv-builder`-supported env var *and*
+                    // that the env var's contents are fully contained in the
+                    // `-C llvm-args` this `rustc` invocation is seeing.
+                    args.windows(args_from_env.len())
+                        .enumerate()
+                        .filter(|&(_, w)| w == args_from_env)
+                        .map(|(w_start, w)| w_start..w_start + w.len())
+                        .flat_map(|w_range| help_flag_positions.range(w_range))
+                        .next()
+                })
+                .is_some();
+            let codegen_args_lhs = if help_flag_comes_from_spirv_builder_env_var {
+                spirv_builder_env_var
+            } else {
+                "rustc -Cllvm-args"
+            };
+            println!(
+                "{}",
+                opts.usage(&format!(
+                    "Usage: {codegen_args_lhs}=\"...\" with `...` from:"
+                ))
+            );
+            // HACK(eddyb) this avoids `Cargo` continuing after the message is printed.
+            std::process::exit(1);
+        }
+
         let module_output_type =
             matches.opt_get_default("module-output", ModuleOutputType::Single)?;
         let disassemble = matches.opt_present("disassemble");
@@ -312,13 +467,19 @@ impl CodegenArgs {
 
         let spirv_metadata = matches.opt_str("spirv-metadata");
 
+        // FIXME(eddyb) clean up this `no-` "negation prefix" situation.
+        let run_spirv_val = !matches.opt_present("no-spirv-val");
+
         let relax_struct_store = matches.opt_present("relax-struct-store");
         let relax_logical_pointer = matches.opt_present("relax-logical-pointer");
         let relax_block_layout = matches.opt_present("relax-block-layout");
         let uniform_buffer_standard_layout = matches.opt_present("uniform-buffer-standard-layout");
         let scalar_block_layout = matches.opt_present("scalar-block-layout");
         let skip_block_layout = matches.opt_present("skip-block-layout");
 
+        // FIXME(eddyb) clean up this `no-` "negation prefix" situation.
+        let run_spirv_opt = !matches.opt_present("no-spirv-opt");
+
         let preserve_bindings = matches.opt_present("preserve-bindings");
 
         let relax_block_layout = if relax_block_layout { Some(true) } else { None };
@@ -334,6 +495,37 @@ impl CodegenArgs {
             }
         };
 
+        let matches_opt_path = |name| matches.opt_str(name).map(PathBuf::from);
+        let matches_opt_dump_dir_path = |name| {
+            matches_opt_path(name).map(|path| {
+                if path.is_file() {
+                    std::fs::remove_file(&path).unwrap();
+                }
+                std::fs::create_dir_all(&path).unwrap();
+                path
+            })
+        };
+        // FIXME(eddyb) should these be handled as `-C linker-args="..."` instead?
+        let linker_opts = crate::linker::Options {
+            // FIXME(eddyb) clean up this `no-` "negation prefix" situation.
+            dce: !matches.opt_present("no-dce"),
+            compact_ids: !matches.opt_present("no-compact-ids"),
+            structurize: !matches.opt_present("no-structurize"),
+
+            // FIXME(eddyb) deduplicate between `CodegenArgs` and `linker::Options`.
+            emit_multiple_modules: module_output_type == ModuleOutputType::Multiple,
+            spirv_metadata,
+
+            // NOTE(eddyb) these are debugging options that used to be env vars
+            // (for more information see `docs/src/codegen-args.md`).
+            dump_post_merge: matches_opt_path("dump-post-merge"),
+            dump_post_split: matches_opt_dump_dir_path("dump-post-split"),
+            specializer_debug: matches.opt_present("specializer-debug"),
+            specializer_dump_instances: matches_opt_path("specializer-dump-instances"),
+            print_all_zombie: matches.opt_present("print-all-zombie"),
+            print_zombie: matches.opt_present("print-zombie"),
+        };
+
         Ok(Self {
             module_output_type,
             disassemble,
@@ -343,14 +535,27 @@ impl CodegenArgs {
 
             spirv_metadata,
 
+            run_spirv_val,
+
             relax_struct_store,
             relax_logical_pointer,
             relax_block_layout,
             uniform_buffer_standard_layout,
             scalar_block_layout,
             skip_block_layout,
 
+            run_spirv_opt,
+
             preserve_bindings,
+
+            linker_opts,
+
+            // NOTE(eddyb) these are debugging options that used to be env vars
+            // (for more information see `docs/src/codegen-args.md`).
+            dump_mir: matches_opt_dump_dir_path("dump-mir"),
+            dump_module_on_panic: matches_opt_path("dump-module-on-panic"),
+            dump_pre_link: matches_opt_dump_dir_path("dump-pre-link"),
+            dump_post_link: matches_opt_dump_dir_path("dump-post-link"),
         })
     }