Clean up DUMP environment variables (#673)

crates/rustc_codegen_spirv/src/lib.rs

@@ -546,8 +546,7 @@ impl ExtraBackendMethods for SpirvCodegenBackend {
         let do_codegen = || {
             let mono_items = cx.codegen_unit.items_in_deterministic_order(cx.tcx);
 
-            if let Ok(path) = env::var("DUMP_MIR") {
-                let mut path = PathBuf::from(path);
+            if let Some(mut path) = get_env_dump_dir("DUMP_MIR") {
                 path.push(cgu_name.to_string());
                 dump_mir(tcx, &mono_items, &path);
             }
@@ -630,6 +629,19 @@ impl Drop for DumpModuleOnPanic<'_, '_, '_> {
     }
 }
 
+fn get_env_dump_dir(env_var: &str) -> Option<PathBuf> {
+    if let Some(path) = std::env::var_os(env_var) {
+        let path = PathBuf::from(path);
+        if path.is_file() {
+            std::fs::remove_file(&path).unwrap();
+        }
+        std::fs::create_dir_all(&path).unwrap();
+        Some(path)
+    } else {
+        None
+    }
+}
+
 /// This is the entrypoint for a hot plugged `rustc_codegen_spirv`
 #[no_mangle]
 pub fn __rustc_codegen_backend() -> Box<dyn CodegenBackend> {

crates/rustc_codegen_spirv/src/link.rs

@@ -2,6 +2,7 @@ use crate::codegen_cx::{CodegenArgs, ModuleOutputType};
 use crate::{
     linker, CompileResult, ModuleResult, SpirvCodegenBackend, SpirvModuleBuffer, SpirvThinBuffer,
 };
+use rspirv::binary::Assemble;
 use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule, ThinShared};
 use rustc_codegen_ssa::back::write::CodegenContext;
 use rustc_codegen_ssa::{CodegenResults, NativeLib};
@@ -21,7 +22,7 @@ use rustc_span::symbol::Symbol;
 use std::env;
 use std::ffi::{CString, OsStr};
 use std::fs::File;
-use std::io::{BufWriter, Read, Write};
+use std::io::{BufWriter, Read};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use tar::{Archive, Builder, Header};
@@ -141,7 +142,6 @@ fn link_exe(
         std::fs::create_dir_all(&out_dir).unwrap();
     }
 
-    use rspirv::binary::Assemble;
     let compile_result = match spv_binary {
         linker::LinkResult::SingleModule(spv_binary) => {
             let mut module_filename = out_dir;
@@ -190,11 +190,9 @@ fn post_link_single_module(
     spv_binary: Vec<u32>,
     out_filename: &Path,
 ) {
-    if let Ok(ref path) = std::env::var("DUMP_POST_LINK") {
-        File::create(path)
-            .unwrap()
-            .write_all(spirv_tools::binary::from_binary(&spv_binary))
-            .unwrap();
+    if let Some(mut path) = crate::get_env_dump_dir("DUMP_POST_LINK") {
+        path.push(out_filename.file_name().unwrap());
+        std::fs::write(path, spirv_tools::binary::from_binary(&spv_binary)).unwrap();
     }
 
     let val_options = spirv_tools::val::ValidatorOptions {
@@ -217,8 +215,21 @@ fn post_link_single_module(
     let spv_binary = if sess.opts.optimize != OptLevel::No
         || (sess.opts.debuginfo == DebugInfo::None && !cg_args.name_variables)
     {
-        let _timer = sess.timer("link_spirv_opt");
-        do_spirv_opt(sess, cg_args, spv_binary, out_filename, opt_options)
+        if env::var("NO_SPIRV_OPT").is_ok() {
+            let _timer = sess.timer("link_spirv_opt");
+            do_spirv_opt(sess, cg_args, spv_binary, out_filename, opt_options)
+        } else {
+            let reason = match (sess.opts.optimize, sess.opts.debuginfo == DebugInfo::None) {
+                (OptLevel::No, true) => "debuginfo=None".to_string(),
+                (optlevel, false) => format!("optlevel={:?}", optlevel),
+                (optlevel, true) => format!("optlevel={:?}, debuginfo=None", optlevel),
+            };
+            sess.warn(&format!(
+                "spirv-opt should have ran ({}) but was disabled by NO_SPIRV_OPT",
+                reason
+            ));
+            spv_binary
+        }
     } else {
         spv_binary
     };
@@ -511,18 +522,13 @@ fn do_link(
         }
     }
 
-    if let Ok(ref path) = env::var("DUMP_PRE_LINK") {
-        use rspirv::binary::Assemble;
-        let path = Path::new(path);
-        if path.is_file() {
-            std::fs::remove_file(path).unwrap();
-        }
-        std::fs::create_dir_all(path).unwrap();
+    if let Some(dir) = crate::get_env_dump_dir("DUMP_PRE_LINK") {
         for (num, module) in modules.iter().enumerate() {
-            File::create(path.join(format!("mod_{}.spv", num)))
-                .unwrap()
-                .write_all(spirv_tools::binary::from_binary(&module.assemble()))
-                .unwrap();
+            std::fs::write(
+                dir.join(format!("mod_{}.spv", num)),
+                spirv_tools::binary::from_binary(&module.assemble()),
+            )
+            .unwrap();
         }
     }
     drop(load_modules_timer);

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -13,7 +13,7 @@ mod structurizer;
 mod zombies;
 
 use crate::decorations::{CustomDecoration, UnrollLoopsDecoration};
-use rspirv::binary::Consumer;
+use rspirv::binary::{Assemble, Consumer};
 use rspirv::dr::{Block, Instruction, Loader, Module, ModuleHeader, Operand};
 use rspirv::spirv::{Op, StorageClass, Word};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
@@ -111,14 +111,7 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
     };
 
     if let Ok(ref path) = std::env::var("DUMP_POST_MERGE") {
-        use rspirv::binary::Assemble;
-        use std::fs::File;
-        use std::io::Write;
-
-        File::create(path)
-            .unwrap()
-            .write_all(spirv_tools::binary::from_binary(&output.assemble()))
-            .unwrap();
+        std::fs::write(path, spirv_tools::binary::from_binary(&output.assemble())).unwrap();
     }
 
     // remove duplicates (https://github.com/KhronosGroup/SPIRV-Tools/blob/e7866de4b1dc2a7e8672867caeb0bdca49f458d3/source/opt/remove_duplicates_pass.cpp)
@@ -283,16 +276,10 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
         LinkResult::SingleModule(ref mut m) => Box::new(std::iter::once(m)),
         LinkResult::MultipleModules(ref mut m) => Box::new(m.values_mut()),
     };
-    for output in output_module_iter {
-        if let Ok(ref path) = std::env::var("DUMP_POST_SPLIT") {
-            use rspirv::binary::Assemble;
-            use std::fs::File;
-            use std::io::Write;
-
-            File::create(path)
-                .unwrap()
-                .write_all(spirv_tools::binary::from_binary(&output.assemble()))
-                .unwrap();
+    for (i, output) in output_module_iter.enumerate() {
+        if let Some(mut path) = crate::get_env_dump_dir("DUMP_POST_SPLIT") {
+            path.push(format!("mod_{}.spv", i));
+            std::fs::write(path, spirv_tools::binary::from_binary(&output.assemble())).unwrap();
         }
         // Run DCE again, even if emit_multiple_modules==false - the first DCE ran before
         // structurization and mem2reg (for perf reasons), and mem2reg may remove references to

docs/src/SUMMARY.md

@@ -4,6 +4,7 @@
 - [Contributing to Rust-GPU]()
     - [Building](./building-rust-gpu.md)
     - [Testing](./testing.md)
+    - [Environment variables Rust-GPU reads](./compiler-env-vars.md)
     - [Minimizing bugs in SPIR-V](./spirv-minimization.md)
 - [Platform Support](./platform-support.md)
 - [Writing Shader Crates](./writing-shader-crates.md)

docs/src/compiler-env-vars.md

@@ -0,0 +1,117 @@
+# Debug environment variables that the Rust-GPU compiler reads
+
+Please keep in mind that all of these variables are for internal development, and may break output
+unexpectedly and generally muck things up. Please only use these if you know what you're doing.
+
+Help is also appreciated keeping this document up to date, environment variables may be
+added/removed on an ad-hoc basis without much thought, as they're internal development tools, not a
+public API - this documentation is only here because these variables may be helpful diagnosing
+problems for others.
+
+It's recommended that environment variables that take paths to files or directories are set to full
+paths, as the working directory of the compiler might be something wonky and unexpected, and it's
+easier to set the full path.
+
+## DUMP_MIR
+
+Takes: path to file
+
+Dumps the MIR of every function rust-gpu encounters to a file. Yes, rustc does have options to do
+this by default, but I always forget the syntax, and plumbing through the option to spirv-builder
+is annoying, so this is handy to just hack an output.
+
+## DUMP_MODULE_ON_PANIC
+
+Takes: path to file
+
+If codegen panics, then write the (partially) emitted module to a file. Note that this only exists
+for codegen, if the linker panics, this option does nothing, sadly.
+
+## DUMP_PRE_LINK
+
+Takes: path to directory
+
+Dumps all input modules to the linker, before the linker touches them at all.
+
+## DUMP_POST_MERGE
+
+Takes: path to file
+
+Dumps the merged module immediately after merging, but before the linker has done anything else
+(including, well, linking the methods - LinkageAttributes will still exist, etc.). This is very
+similar to DUMP_PRE_LINK, except it outputs only a single file, which might make grepping through
+for stuff easier.
+
+## DUMP_POST_SPLIT
+
+Takes: path to directory
+
+Dumps the modules immediately after multimodule splitting, but before final cleanup passes (e.g.
+DCE to remove the other entry points).
+
+## DUMP_POST_LINK
+
+Takes: path to directory
+
+Dumps all output modules from the linker. This may be multiple files due to the multimodule/module
+splitting option, hence it takes a directory instead of a file path. This is the final output
+binary before spirv-opt is executed, so it may be useful to output this to check if an issue is in
+rust-gpu, or in spirv-opt.
+
+## SPECIALIZER_DEBUG
+
+Takes: presence or absence (e.g. set to `1`)
+
+Prints to rustc stdout some debug information about the specializer.
+
+## SPECIALIZER_DUMP_INSTANCES
+
+Takes: path to file
+
+Prints to file some... stuff, idk, ask eddyb (useful for debugging specializer)
+
+## PRINT_ZOMBIE
+
+Takes: presence or absence (e.g. set to `1`)
+
+Prints to rustc stdout which functions were removed due to being zombies, and why.
+
+## PRINT_ALL_ZOMBIE
+
+Takes: presence or absence (e.g. set to `1`)
+
+Prints to rustc stdout *everything* that was removed due to being zombies, why, and if it was an
+original zombie or if it was infected. (prints a lot!)
+
+## NO_SPIRV_VAL
+
+Takes: presence or absence (e.g. set to `1`)
+
+Disables running spirv-val on the final output. Spooky scary option, can cause invalid modules!
+
+## NO_SPIRV_OPT
+
+Takes: presence or absence (e.g. set to `1`)
+
+Forcibly disables running spirv-opt on the final output, even if optimizations are enabled.
+
+## NO_DCE
+
+Takes: presence or absence (e.g. set to `1`)
+
+Disables running dead code elimination. Can and probably will generate invalid modules or crash the
+linker, hasn't been tested for a while.
+
+## NO_COMPACT_IDS
+
+Takes: presence or absence (e.g. set to `1`)
+
+Disables compaction of SPIR-V IDs at the end of linking. Causes absolutely ginormous IDs to be
+emitted. Useful if you're println debugging IDs in the linker (although spirv-opt will compact them
+anyway, be careful).
+
+## NO_STRUCTURIZE
+
+Takes: presence or absence (e.g. set to `1`)
+
+Disables structurization. Probably results in invalid modules.