Include identifying information (e.g. crate names) in --dump-* file…

…names.
EmbarkStudios · Dec 13, 2022 · b1b1190 · b1b1190
1 parent 6fb35ff
commit b1b1190
Show file tree

Hide file tree

Showing 9 changed files with 201 additions and 91 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -33,14 +33,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - [PR#959](https://github.com/EmbarkStudios/rust-gpu/pull/959) 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)
+    (check out [the "codegen args" docs](docs/src/codegen-args.md), or run with `RUSTGPU_CODEGEN_ARGS=--help` to see the full list of options)
 - [PR#940](https://github.com/EmbarkStudios/rust-gpu/pull/940) integrated the experimental [`SPIR-🇹` shader IR framework](https://github.com/EmbarkStudios/spirt) into the linker  
   (opt-in via `RUSTGPU_CODEGEN_ARGS=--spirt`, see also [the `--spirt` docs](docs/src/codegen-args.md#--spirt), for more details)
 
 ### Changed 🛠️
 - [PR#958](https://github.com/EmbarkStudios/rust-gpu/pull/958) updated toolchain to `nightly-2022-10-29`
 - [PR#941](https://github.com/EmbarkStudios/rust-gpu/pull/941) applied workspace inheritance to `Cargo.toml` files
 - [PR#959](https://github.com/EmbarkStudios/rust-gpu/pull/959) 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)
+- [PR#967](https://github.com/EmbarkStudios/rust-gpu/pull/967) made `--dump-*` ["codegen args"](docs/src/codegen-args.md) include identifying information (e.g. crate names) in the names of files they emit
 
 ### Removed 🔥
 - [PR#946](https://github.com/EmbarkStudios/rust-gpu/pull/946) removed the `fn`/closure `#[spirv(unroll_loops)]` attribute, as it has no users, is becoming non-trivial to support, and requires redesign for better ergonomics (e.g. `#[spirv(unroll)]` applied to individual loops, not the whole `fn`/closure)

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/rustc_codegen_spirv-types/src/compile_result.rs b/crates/rustc_codegen_spirv-types/src/compile_result.rs
@@ -32,8 +32,8 @@ impl ModuleResult {
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct CompileResult {
-    pub module: ModuleResult,
     pub entry_points: Vec<String>,
+    pub module: ModuleResult,
 }
 
 impl CompileResult {

diff --git a/crates/rustc_codegen_spirv/Cargo.toml b/crates/rustc_codegen_spirv/Cargo.toml
@@ -47,6 +47,7 @@ regex = { version = "1", features = ["perf"] }
 
 # Normal dependencies.
 ar = "0.9.0"
+either = "1.8.0"
 indexmap = "1.6.0"
 rspirv = "0.11"
 rustc-demangle = "0.1.21"

diff --git a/crates/rustc_codegen_spirv/src/codegen_cx/mod.rs b/crates/rustc_codegen_spirv/src/codegen_cx/mod.rs
@@ -362,8 +362,8 @@ impl CodegenArgs {
             opts.optopt(
                 "",
                 "dump-post-merge",
-                "dump the merged module immediately after merging, to FILE",
-                "FILE",
+                "dump the merged module immediately after merging, to a file in DIR",
+                "DIR",
             );
             opts.optopt(
                 "",
@@ -374,8 +374,8 @@ impl CodegenArgs {
             opts.optopt(
                 "",
                 "dump-spirt-passes",
-                "dump the SPIR-T module across passes, to FILE and FILE.html",
-                "FILE",
+                "dump the SPIR-T module across passes, to a (pair of) file(s) in DIR",
+                "DIR",
             );
             opts.optflag(
                 "",
@@ -532,9 +532,9 @@ impl CodegenArgs {
 
             // 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_merge: matches_opt_dump_dir_path("dump-post-merge"),
             dump_post_split: matches_opt_dump_dir_path("dump-post-split"),
-            dump_spirt_passes: matches_opt_path("dump-spirt-passes"),
+            dump_spirt_passes: matches_opt_dump_dir_path("dump-spirt-passes"),
             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"),

diff --git a/crates/rustc_codegen_spirv/src/link.rs b/crates/rustc_codegen_spirv/src/link.rs
@@ -2,6 +2,7 @@ use crate::codegen_cx::{CodegenArgs, SpirvMetadata};
 use crate::{linker, SpirvCodegenBackend, SpirvModuleBuffer, SpirvThinBuffer};
 use ar::{Archive, GnuBuilder, Header};
 use rspirv::binary::Assemble;
+use rspirv::dr::Module;
 use rustc_ast::CRATE_NODE_ID;
 use rustc_codegen_spirv_types::{CompileResult, ModuleResult};
 use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule, ThinShared};
@@ -17,7 +18,8 @@ use rustc_session::config::{CrateType, DebugInfo, Lto, OptLevel, OutputFilenames
 use rustc_session::output::{check_file_is_writeable, invalid_output_for_target, out_filename};
 use rustc_session::utils::NativeLibKind;
 use rustc_session::Session;
-use std::ffi::CString;
+use std::collections::BTreeMap;
+use std::ffi::{CString, OsStr};
 use std::fs::File;
 use std::io::{BufWriter, Read};
 use std::iter;
@@ -62,7 +64,21 @@ pub fn link<'a>(
                     link_rlib(sess, codegen_results, &out_filename);
                 }
                 CrateType::Executable | CrateType::Cdylib | CrateType::Dylib => {
-                    link_exe(sess, crate_type, &out_filename, codegen_results);
+                    // HACK(eddyb) there's no way way to access `outputs.filestem`,
+                    // so we pay the cost of building a whole `PathBuf` instead.
+                    let disambiguated_crate_name_for_dumps = outputs
+                        .with_extension("")
+                        .file_name()
+                        .unwrap()
+                        .to_os_string();
+
+                    link_exe(
+                        sess,
+                        crate_type,
+                        &out_filename,
+                        codegen_results,
+                        &disambiguated_crate_name_for_dumps,
+                    );
                 }
                 other => {
                     sess.err(format!("CrateType {:?} not supported yet", other));
@@ -117,6 +133,7 @@ fn link_exe(
     crate_type: CrateType,
     out_filename: &Path,
     codegen_results: &CodegenResults,
+    disambiguated_crate_name_for_dumps: &OsStr,
 ) {
     let mut objects = Vec::new();
     let mut rlibs = Vec::new();
@@ -140,38 +157,49 @@ fn link_exe(
     // HACK(eddyb) this removes the `.json` in `.spv.json`, from `out_filename`.
     let out_path_spv = out_filename.with_extension("");
 
-    let compile_result = match do_link(sess, &cg_args, &objects, &rlibs) {
+    let link_result = do_link(
+        sess,
+        &cg_args,
+        &objects,
+        &rlibs,
+        disambiguated_crate_name_for_dumps,
+    );
+    let compile_result = match link_result {
         linker::LinkResult::SingleModule(module) => {
-            let module_filename = out_path_spv;
-            post_link_single_module(sess, &cg_args, module.assemble(), &module_filename);
-            cg_args.do_disassemble(&module);
-            let module_result = ModuleResult::SingleModule(module_filename);
+            let entry_points = entry_points(&module);
+            post_link_single_module(sess, &cg_args, *module, &out_path_spv, None);
             CompileResult {
-                module: module_result,
-                entry_points: entry_points(&module),
+                entry_points,
+                module: ModuleResult::SingleModule(out_path_spv),
             }
         }
-        linker::LinkResult::MultipleModules(map) => {
+        linker::LinkResult::MultipleModules {
+            file_stem_to_entry_name_and_module,
+        } => {
             let out_dir = out_path_spv.with_extension("spvs");
             if !out_dir.is_dir() {
                 std::fs::create_dir_all(&out_dir).unwrap();
             }
 
-            let entry_points = map.keys().cloned().collect();
-            let map = map
+            let entry_name_to_file_path: BTreeMap<_, _> = file_stem_to_entry_name_and_module
                 .into_iter()
-                .map(|(name, module)| {
-                    let mut module_filename = out_dir.clone();
-                    module_filename.push(sanitize_filename::sanitize(&name));
-                    module_filename.set_extension("spv");
-                    post_link_single_module(sess, &cg_args, module.assemble(), &module_filename);
-                    (name, module_filename)
+                .map(|(file_stem, (entry_name, module))| {
+                    let mut out_file_name = file_stem;
+                    out_file_name.push(".spv");
+                    let out_file_path = out_dir.join(out_file_name);
+                    post_link_single_module(
+                        sess,
+                        &cg_args,
+                        module,
+                        &out_file_path,
+                        Some(disambiguated_crate_name_for_dumps),
+                    );
+                    (entry_name, out_file_path)
                 })
                 .collect();
-            let module_result = ModuleResult::MultiModule(map);
             CompileResult {
-                module: module_result,
-                entry_points,
+                entry_points: entry_name_to_file_path.keys().cloned().collect(),
+                module: ModuleResult::MultiModule(entry_name_to_file_path),
             }
         }
     };
@@ -192,15 +220,21 @@ fn entry_points(module: &rspirv::dr::Module) -> Vec<String> {
 fn post_link_single_module(
     sess: &Session,
     cg_args: &CodegenArgs,
-    spv_binary: Vec<u32>,
+    module: Module,
     out_filename: &Path,
+    dump_prefix: Option<&OsStr>,
 ) {
+    cg_args.do_disassemble(&module);
+    let spv_binary = module.assemble();
+
     if let Some(dir) = &cg_args.dump_post_link {
-        std::fs::write(
-            dir.join(out_filename.file_name().unwrap()),
-            spirv_tools::binary::from_binary(&spv_binary),
-        )
-        .unwrap();
+        // FIXME(eddyb) rename `filename` with `file_path` to make this less confusing.
+        let out_filename_file_name = out_filename.file_name().unwrap();
+        let dump_path = match dump_prefix {
+            Some(prefix) => dir.join(prefix).with_extension(out_filename_file_name),
+            None => dir.join(out_filename_file_name),
+        };
+        std::fs::write(dump_path, spirv_tools::binary::from_binary(&spv_binary)).unwrap();
     }
 
     let val_options = spirv_tools::val::ValidatorOptions {
@@ -514,20 +548,34 @@ fn do_link(
     cg_args: &CodegenArgs,
     objects: &[PathBuf],
     rlibs: &[PathBuf],
+    disambiguated_crate_name_for_dumps: &OsStr,
 ) -> linker::LinkResult {
-    fn load(bytes: &[u8]) -> rspirv::dr::Module {
-        let mut loader = rspirv::dr::Loader::new();
-        rspirv::binary::parse_bytes(bytes, &mut loader).unwrap();
-        loader.module()
-    }
-
     let load_modules_timer = sess.timer("link_load_modules");
+
     let mut modules = Vec::new();
+    let mut add_module = |file_name: &OsStr, bytes: &[u8]| {
+        let module = {
+            let mut loader = rspirv::dr::Loader::new();
+            rspirv::binary::parse_bytes(bytes, &mut loader).unwrap();
+            loader.module()
+        };
+        if let Some(dir) = &cg_args.dump_pre_link {
+            // FIXME(eddyb) is it a good idea to re-`assemble` the `rspirv::dr`
+            // module, or should this just save the original bytes?
+            std::fs::write(
+                dir.join(file_name).with_extension("spv"),
+                spirv_tools::binary::from_binary(&module.assemble()),
+            )
+            .unwrap();
+        }
+        modules.push(module);
+    };
+
     // `objects` are the plain obj files we need to link - usually produced by the final crate.
     for obj in objects {
-        let bytes = std::fs::read(obj).unwrap();
-        modules.push(load(&bytes));
+        add_module(obj.file_name().unwrap(), &std::fs::read(obj).unwrap());
     }
+
     // `rlibs` are archive files we've created in `create_archive`, usually produced by crates that are being
     // referenced. We need to unpack them and add the modules inside.
     for rlib in rlibs {
@@ -539,24 +587,22 @@ fn do_link(
                 // https://github.com/rust-lang/rust/blob/72868e017bdade60603a25889e253f556305f996/library/std/src/fs.rs#L200-L202
                 let mut bytes = Vec::with_capacity(entry.header().size() as usize + 1);
                 entry.read_to_end(&mut bytes).unwrap();
-                modules.push(load(&bytes));
+
+                let file_name = std::str::from_utf8(entry.header().identifier()).unwrap();
+                add_module(OsStr::new(file_name), &bytes);
             }
         }
     }
 
-    if let Some(dir) = &cg_args.dump_pre_link {
-        for (num, module) in modules.iter().enumerate() {
-            std::fs::write(
-                dir.join(format!("mod_{}.spv", num)),
-                spirv_tools::binary::from_binary(&module.assemble()),
-            )
-            .unwrap();
-        }
-    }
     drop(load_modules_timer);
 
     // Do the link...
-    let link_result = linker::link(sess, modules, &cg_args.linker_opts);
+    let link_result = linker::link(
+        sess,
+        modules,
+        &cg_args.linker_opts,
+        disambiguated_crate_name_for_dumps,
+    );
 
     match link_result {
         Ok(v) => v,