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,
+                    }
+                }
+            }
+        }
+    }
+}