linker: add --spirt-passes codegen args and underlying abstraction.

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -364,6 +364,12 @@ impl CodegenArgs {
                 "spirt",
                 "use SPIR-T for legalization (see also `docs/src/codegen-args.md`)",
             );
+            opts.optmulti(
+                "",
+                "spirt-passes",
+                "enable additional SPIR-T passes (comma-separated)",
+                "PASSES",
+            );
 
             // NOTE(eddyb) these are debugging options that used to be env vars
             // (for more information see `docs/src/codegen-args.md`).
@@ -532,6 +538,12 @@ impl CodegenArgs {
             compact_ids: !matches.opt_present("no-compact-ids"),
             structurize: !matches.opt_present("no-structurize"),
             spirt: matches.opt_present("spirt"),
+            spirt_passes: matches
+                .opt_strs("spirt-passes")
+                .iter()
+                .flat_map(|s| s.split(','))
+                .map(|s| s.to_string())
+                .collect(),
 
             // FIXME(eddyb) deduplicate between `CodegenArgs` and `linker::Options`.
             emit_multiple_modules: module_output_type == ModuleOutputType::Multiple,

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -13,6 +13,7 @@ mod param_weakening;
 mod peephole_opts;
 mod simple_passes;
 mod specializer;
+mod spirt_passes;
 mod structurizer;
 mod zombies;
 
@@ -38,6 +39,7 @@ pub struct Options {
     pub dce: bool,
     pub structurize: bool,
     pub spirt: bool,
+    pub spirt_passes: Vec<String>,
 
     pub emit_multiple_modules: bool,
     pub spirv_metadata: SpirvMetadata,
@@ -396,6 +398,18 @@ pub fn link(
             after_pass("structurize_func_cfgs", &module);
         }
 
+        if !opts.spirt_passes.is_empty() {
+            spirt_passes::run_func_passes(
+                &mut module,
+                &opts.spirt_passes,
+                |name, _module| sess.timer(name),
+                |name, module, timer| {
+                    drop(timer);
+                    after_pass(name, module);
+                },
+            );
+        }
+
         // NOTE(eddyb) this should be *before* `lift_to_spv` below,
         // so if that fails, the dump could be used to debug it.
         if let Some(dump_dir) = &opts.dump_spirt_passes {
@@ -413,13 +427,21 @@ pub fn link(
 
             // FIXME(eddyb) don't allocate whole `String`s here.
             std::fs::write(&dump_spirt_file_path, pretty.to_string()).unwrap();
-            std::fs::write(
-                dump_spirt_file_path.with_extension("spirt.html"),
-                pretty
+            std::fs::write(dump_spirt_file_path.with_extension("spirt.html"), {
+                let mut html = pretty
                     .render_to_html()
                     .with_dark_mode_support()
-                    .to_html_doc(),
-            )
+                    .to_html_doc();
+                // HACK(eddyb) this should be in `spirt::pretty` itself,
+                // but its need didn't become obvious until more recently.
+                html += "
+                        <style>
+                            pre.spirt-90c2056d-5b38-4644-824a-b4be1c82f14d sub {
+                                line-height: 0;
+                            }
+                        </style>";
+                html
+            })
             .unwrap();
         }
 

crates/rustc_codegen_spirv/src/linker/spirt_passes/mod.rs

@@ -0,0 +1,88 @@
+//! SPIR-T pass infrastructure and supporting utilities.
+
+use rustc_data_structures::fx::FxIndexSet;
+use spirt::visit::{InnerVisit, Visitor};
+use spirt::{AttrSet, Const, Context, DeclDef, Func, GlobalVar, Module, Type};
+
+/// Run intra-function passes on all `Func` definitions in the `Module`.
+//
+// FIXME(eddyb) introduce a proper "pass manager".
+pub(super) fn run_func_passes<P>(
+    module: &mut Module,
+    passes: &[impl AsRef<str>],
+    // FIXME(eddyb) this is a very poor approximation of a "profiler" abstraction.
+    mut before_pass: impl FnMut(&'static str, &Module) -> P,
+    mut after_pass: impl FnMut(&'static str, &Module, P),
+) {
+    let cx = &module.cx();
+
+    // FIXME(eddyb) reuse this collection work in some kind of "pass manager".
+    let all_funcs = {
+        let mut collector = ReachableUseCollector {
+            cx,
+            module,
+
+            seen_types: FxIndexSet::default(),
+            seen_consts: FxIndexSet::default(),
+            seen_global_vars: FxIndexSet::default(),
+            seen_funcs: FxIndexSet::default(),
+        };
+        for &exportee in module.exports.values() {
+            exportee.inner_visit_with(&mut collector);
+        }
+        collector.seen_funcs
+    };
+
+    for name in passes {
+        let name = name.as_ref();
+        let (full_name, pass_fn): (_, fn(_, &mut _)) = match name {
+            _ => panic!("unknown `--spirt-passes={}`", name),
+        };
+
+        let profiler = before_pass(full_name, module);
+        for &func in &all_funcs {
+            if let DeclDef::Present(func_def_body) = &mut module.funcs[func].def {
+                pass_fn(cx, func_def_body);
+            }
+        }
+        after_pass(full_name, module, profiler);
+    }
+}
+
+// FIXME(eddyb) this is just copy-pasted from `spirt` and should be reusable.
+struct ReachableUseCollector<'a> {
+    cx: &'a Context,
+    module: &'a Module,
+
+    // FIXME(eddyb) build some automation to avoid ever repeating these.
+    seen_types: FxIndexSet<Type>,
+    seen_consts: FxIndexSet<Const>,
+    seen_global_vars: FxIndexSet<GlobalVar>,
+    seen_funcs: FxIndexSet<Func>,
+}
+
+impl Visitor<'_> for ReachableUseCollector<'_> {
+    // FIXME(eddyb) build some automation to avoid ever repeating these.
+    fn visit_attr_set_use(&mut self, _attrs: AttrSet) {}
+    fn visit_type_use(&mut self, ty: Type) {
+        if self.seen_types.insert(ty) {
+            self.visit_type_def(&self.cx[ty]);
+        }
+    }
+    fn visit_const_use(&mut self, ct: Const) {
+        if self.seen_consts.insert(ct) {
+            self.visit_const_def(&self.cx[ct]);
+        }
+    }
+
+    fn visit_global_var_use(&mut self, gv: GlobalVar) {
+        if self.seen_global_vars.insert(gv) {
+            self.visit_global_var_decl(&self.module.global_vars[gv]);
+        }
+    }
+    fn visit_func_use(&mut self, func: Func) {
+        if self.seen_funcs.insert(func) {
+            self.visit_func_decl(&self.module.funcs[func]);
+        }
+    }
+}

docs/src/codegen-args.md

@@ -149,6 +149,12 @@ Enables using the experimental [`SPIR-🇹` shader IR framework](https://github.
 
 For more information, also see [the `SPIR-🇹` repository](https://github.com/EmbarkStudios/spirt).
 
+### `--spirt-passes PASSES`
+
+Enable additional `SPIR-🇹` passes, as listed in `PASSES` (comma-separated).
+_Note: passes that are not already enabled by default are considered experimental
+and likely not ready for production use, this flag exists primarily for testing.*
+
 ### `--dump-spirt-passes DIR`
 
 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`.