Add --spirt-passes and a reduce pass for replacing ops with their inputs/constants.

CHANGELOG.md

@@ -29,6 +29,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added ⭐
+- [PR#988](https://github.com/EmbarkStudios/rust-gpu/pull/988) added a couple of (highly experimental)
+ `SPIR-🇹` optimization passes, and `--spirt-passes=...` codegen args as a way to enable them 
+ (see also [the `--spirt-passes` codegen args docs](docs/src/codegen-args.md#--spirt-passes-PASSES))
+
 ### Changed 🛠️
 - [PR#982](https://github.com/EmbarkStudios/rust-gpu/pull/982) updated toolchain to `nightly-2022-12-18`
 - [PR#953](https://github.com/EmbarkStudios/rust-gpu/pull/953) migrated to the Rust 2021 edition, and fixed Rust 2021 support for shader crates to be on par with Rust 2018 (discrepancies having been limited to/caused by `panic!` changes in Rust 2021)

crates/rustc_codegen_spirv/Cargo.toml

@@ -59,6 +59,7 @@ smallvec = { version = "1.6.1", features = ["union"] }
 spirv-tools = { version = "0.9", default-features = false }
 rustc_codegen_spirv-types.workspace = true
 spirt = "0.1.0"
+lazy_static = "1.4.0"
 
 [dev-dependencies]
 pipe = "0.4"

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/fuse_selects.rs

@@ -0,0 +1,108 @@
+use spirt::func_at::FuncAt;
+use spirt::transform::InnerInPlaceTransform;
+use spirt::visit::InnerVisit;
+use spirt::{Context, ControlNodeKind, ControlRegion, FuncDefBody, SelectionKind, Value};
+use std::mem;
+
+use super::{ReplaceValueWith, VisitAllControlRegionsAndNodes};
+
+/// Combine consecutive `Select`s in `func_def_body`.
+pub(crate) fn fuse_selects_in_func(_cx: &Context, func_def_body: &mut FuncDefBody) {
+ // HACK(eddyb) this kind of random-access is easier than using `spirt::transform`.
+ let mut all_regions = vec![];
+
+ func_def_body.inner_visit_with(&mut VisitAllControlRegionsAndNodes {
+ state: (),
+ visit_control_region: |_: &mut (), func_at_control_region: FuncAt<'_, ControlRegion>| {
+ all_regions.push(func_at_control_region.position);
+ },
+ visit_control_node: |_: &mut (), _| {},
+ });
+
+ for region in all_regions {
+ let mut func_at_children_iter = func_def_body.at_mut(region).at_children().into_iter();
+ while let Some(func_at_child) = func_at_children_iter.next() {
+ let base_control_node = func_at_child.position;
+ if let ControlNodeKind::Select {
+ kind: SelectionKind::BoolCond,
+ scrutinee,
+ cases,
+ } = &func_at_child.def().kind
+ {
+ let &base_cond = scrutinee;
+ let base_cases = cases.clone();
+
+ // Scan ahead for candidate `Select`s (with the same condition).
+ let mut fusion_candidate_iter = func_at_children_iter.reborrow();
+ while let Some(func_at_fusion_candidate) = fusion_candidate_iter.next() {
+ let fusion_candidate = func_at_fusion_candidate.position;
+ let mut func = func_at_fusion_candidate.at(());
+ let fusion_candidate_def = func.reborrow().at(fusion_candidate).def();
+ match &fusion_candidate_def.kind {
+ // HACK(eddyb) ignore empty blocks (created by
+ // e.g. `remove_unused_values_in_func`).
+ ControlNodeKind::Block { insts } if insts.is_empty() => {}
+
+ ControlNodeKind::Select {
+ kind: SelectionKind::BoolCond,
+ scrutinee: fusion_candidate_cond,
+ cases: fusion_candidate_cases,
+ } if *fusion_candidate_cond == base_cond => {
+ // FIXME(eddyb) handle outputs from the second `Select`.
+ if !fusion_candidate_def.outputs.is_empty() {
+ break;
+ }
+
+ let cases_to_fuse = fusion_candidate_cases.clone();
+
+ // Concatenate the `Select`s' respective cases
+ // ("then" with "then", "else" with "else", etc.).
+ for (&base_case, &case_to_fuse) in base_cases.iter().zip(&cases_to_fuse)
+ {
+ let children_of_case_to_fuse =
+ mem::take(&mut func.reborrow().at(case_to_fuse).def().children);
+
+ // Replace uses of the outputs of the first `Select`,
+ // in the second one's case, with the specific values
+ // (e.g. `let y = if c { x } ...; if c { f(y) }`
+ // has to become `let y = if c { f(x); x } ...`).
+ //
+ // FIXME(eddyb) avoid cloning here.
+ let outputs_of_base_case =
+ func.reborrow().at(base_case).def().outputs.clone();
+ func.reborrow()
+ .at(children_of_case_to_fuse)
+ .into_iter()
+ .inner_in_place_transform_with(&mut ReplaceValueWith(
+ |v| match v {
+ Value::ControlNodeOutput {
+ control_node,
+ output_idx,
+ } if control_node == base_control_node => {
+ Some(outputs_of_base_case[output_idx as usize])
+ }
+
+ _ => None,
+ },
+ ));
+
+ func.control_regions[base_case]
+ .children
+ .append(children_of_case_to_fuse, func.control_nodes);
+ }
+
+ // HACK(eddyb) because we can't remove list elements yet,
+ // we instead replace the `Select` with an empty `Block`.
+ func.reborrow().at(fusion_candidate).def().kind =
+ ControlNodeKind::Block {
+ insts: Default::default(),
+ };
+ }
+
+ _ => break,
+ }
+ }
+ }
+ }
+ }
+}