Skip to content

Commit

Permalink
linker: add --spirt-passes codegen args and underlying abstraction.
Browse files Browse the repository at this point in the history
  • Loading branch information
eddyb committed Jan 13, 2023
1 parent d2d6ee2 commit 309739d
Show file tree
Hide file tree
Showing 4 changed files with 133 additions and 5 deletions.
12 changes: 12 additions & 0 deletions crates/rustc_codegen_spirv/src/codegen_cx/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -356,6 +356,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`).
Expand Down Expand Up @@ -524,6 +530,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,
Expand Down
32 changes: 27 additions & 5 deletions crates/rustc_codegen_spirv/src/linker/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ mod param_weakening;
mod peephole_opts;
mod simple_passes;
mod specializer;
mod spirt_passes;
mod structurizer;
mod zombies;

Expand All @@ -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,
Expand Down Expand Up @@ -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 {
Expand All @@ -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();
}

Expand Down
88 changes: 88 additions & 0 deletions crates/rustc_codegen_spirv/src/linker/spirt_passes/mod.rs
Original file line number Diff line number Diff line change
@@ -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) = 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]);
}
}
}
6 changes: 6 additions & 0 deletions docs/src/codegen-args.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down

0 comments on commit 309739d

Please sign in to comment.