Add --dump-spirt for dumping only the final SPIR-T module (unlike --dump-spirt-passes).

Author: eddyb

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -331,16 +331,19 @@ pub struct CodegenArgs {
 
 impl CodegenArgs {
     pub fn from_session(sess: &Session) -> Self {
-        match CodegenArgs::parse(&sess.opts.cg.llvm_args) {
+        match Self::parse(sess, &sess.opts.cg.llvm_args) {
             Ok(ok) => ok,
+
+            // FIXME(eddyb) this should mention `RUSTGPU_CODEGEN_ARGS`, just
+            // like how `RUSTGPU_CODEGEN_ARGS=--help` is already special-cased.
             Err(err) => sess
                 .dcx()
                 .fatal(format!("Unable to parse llvm-args: {err}")),
         }
     }
 
-    // FIXME(eddyb) `structopt` would come a long way to making this nicer.
-    pub fn parse(args: &[String]) -> Result<Self, rustc_session::getopts::Fail> {
+    // FIXME(eddyb) switch all of this over to `clap`.
+    pub fn parse(sess: &Session, 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
@@ -459,6 +462,12 @@ impl CodegenArgs {
                 "dump the SPIR-T module across passes, to a (pair of) file(s) in DIR",
                 "DIR",
             );
+            opts.optopt(
+                "",
+                "dump-spirt",
+                "dump the final SPIR-T module, to a (pair of) file(s) in DIR",
+                "DIR",
+            );
             opts.optflag(
                 "",
                 "spirt-strip-custom-debuginfo-from-dumps",
@@ -627,7 +636,19 @@ impl CodegenArgs {
             dump_pre_inline: matches_opt_dump_dir_path("dump-pre-inline"),
             dump_post_inline: matches_opt_dump_dir_path("dump-post-inline"),
             dump_post_split: matches_opt_dump_dir_path("dump-post-split"),
-            dump_spirt_passes: matches_opt_dump_dir_path("dump-spirt-passes"),
+            dump_spirt: match ["dump-spirt-passes", "dump-spirt"].map(matches_opt_dump_dir_path) {
+                [Some(dump_spirt_passes), dump_spirt] => {
+                    if dump_spirt.is_some() {
+                        sess.dcx()
+                            .warn("`--dump-spirt` ignored in favor of `--dump-spirt-passes`");
+                    }
+                    Some((dump_spirt_passes, crate::linker::DumpSpirtMode::AllPasses))
+                }
+                [None, Some(dump_spirt)] => {
+                    Some((dump_spirt, crate::linker::DumpSpirtMode::OnlyFinal))
+                }
+                [None, None] => None,
+            },
             spirt_strip_custom_debuginfo_from_dumps: matches
                 .opt_present("spirt-strip-custom-debuginfo-from-dumps"),
             spirt_keep_debug_sources_in_dumps: matches

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -61,13 +61,19 @@ pub struct Options {
     pub dump_pre_inline: Option<PathBuf>,
     pub dump_post_inline: Option<PathBuf>,
     pub dump_post_split: Option<PathBuf>,
-    pub dump_spirt_passes: Option<PathBuf>,
+    pub dump_spirt: Option<(PathBuf, DumpSpirtMode)>,
     pub spirt_strip_custom_debuginfo_from_dumps: bool,
     pub spirt_keep_debug_sources_in_dumps: bool,
     pub spirt_keep_unstructured_cfg_in_dumps: bool,
     pub specializer_dump_instances: Option<PathBuf>,
 }
 
+#[derive(Copy, Clone, PartialEq, Eq)]
+pub enum DumpSpirtMode {
+    AllPasses,
+    OnlyFinal,
+}
+
 pub enum LinkResult {
     SingleModule(Box<Module>),
     MultipleModules {
@@ -520,7 +526,7 @@ pub fn link(
             drop(timer);
             let pass_name = dump_guard.in_progress_pass_name.take().unwrap();
             if let Some(module) = module
-                && opts.dump_spirt_passes.is_some()
+                && matches!(opts.dump_spirt, Some((_, DumpSpirtMode::AllPasses)))
             {
                 dump_guard
                     .per_pass_module_for_dumping
@@ -801,9 +807,9 @@ impl Drop for SpirtDumpGuard<'_> {
 
         let mut dump_spirt_file_path =
             self.linker_options
-                .dump_spirt_passes
+                .dump_spirt
                 .as_ref()
-                .map(|dump_dir| {
+                .map(|(dump_dir, _)| {
                     dump_dir
                         .join(self.disambiguated_crate_name_for_dumps)
                         .with_extension("spirt")
@@ -815,17 +821,18 @@ impl Drop for SpirtDumpGuard<'_> {
         // but that requires keeping around e.g. the initial SPIR-V for longer,
         // and probably invoking the "SPIR-T pipeline" here, as looping is hard).
         if self.any_spirt_bugs && dump_spirt_file_path.is_none() {
-            if self.per_pass_module_for_dumping.is_empty() {
-                self.per_pass_module_for_dumping
-                    .push(("".into(), self.module.clone()));
-            }
             dump_spirt_file_path = Some(self.outputs.temp_path_for_diagnostic("spirt"));
         }
 
         let Some(dump_spirt_file_path) = &dump_spirt_file_path else {
             return;
         };
 
+        if self.per_pass_module_for_dumping.is_empty() {
+            self.per_pass_module_for_dumping
+                .push(("".into(), self.module.clone()));
+        }
+
         for (_, module) in &mut self.per_pass_module_for_dumping {
             // FIXME(eddyb) consider catching panics in this?
             self.linker_options.spirt_cleanup_for_dumping(module);
@@ -835,7 +842,16 @@ impl Drop for SpirtDumpGuard<'_> {
         let versions = self
             .per_pass_module_for_dumping
             .iter()
-            .map(|(pass_name, module)| (format!("after {pass_name}"), module));
+            .map(|(pass_name, module)| {
+                (
+                    if pass_name.is_empty() {
+                        "".into()
+                    } else {
+                        format!("after {pass_name}")
+                    },
+                    module,
+                )
+            });
 
         let mut panicked_printing_after_passes = None;
         for truncate_version_count in (1..=versions.len()).rev() {
@@ -896,7 +912,11 @@ impl Drop for SpirtDumpGuard<'_> {
                 "pretty-printed SPIR-T was saved to {}.html",
                 dump_spirt_file_path.display()
             ));
-            if self.linker_options.dump_spirt_passes.is_none() {
+            let is_dumping_spirt_passes = matches!(
+                self.linker_options.dump_spirt,
+                Some((_, DumpSpirtMode::AllPasses))
+            );
+            if !is_dumping_spirt_passes {
                 note.help("re-run with `RUSTGPU_CODEGEN_ARGS=\"--dump-spirt-passes=$PWD\"` for more details");
             }
             note.note("pretty-printed SPIR-T is preferred when reporting Rust-GPU issues");

docs/src/codegen-args.md

@@ -175,6 +175,14 @@ _Note: passes that are not already enabled by default are considered experimenta
 Dump the `SPIR-🇹` module across passes (i.e. all of the versions before/after each pass), as a combined report, to a pair of files (`.spirt` and `.spirt.html`) in `DIR`.  
 <sub>(the `.spirt.html` version of the report is the recommended form for viewing, as it uses tabling for versions, syntax-highlighting-like styling, and use->def linking)</sub>
 
+Mutually exclusive with `--dump-spirt` (this takes precedence over that).
+
+### `--dump-spirt DIR`
+
+Dump the `SPIR-🇹` module, similar to `--dump-spirt-passes`, but only the final version.
+
+Mutually exclusive with `--dump-spirt-passes` (which takes precedence over this).
+
 ### `--spirt-strip-custom-debuginfo-from-dumps`
 
 When dumping (pretty-printed) `SPIR-🇹` (e.g. with `--dump-spirt-passes`), strip