Minimally integrate SPIR-T (opt-in via RUSTGPU_CODEGEN_ARGS=--spirt).

Author: eddyb

CHANGELOG.md

@@ -34,6 +34,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - `RUSTGPU_RUSTFLAGS="..."` for shader `RUSTFLAGS="..."`
   - `RUSTGPU_CODEGEN_ARGS="..."` for shader "codegen args" (i.e. `RUSTFLAGS=-Cllvm-args="..."`)  
     (check out ["codegen args" docs](docs/src/codegen-args.md), or run with `RUSTGPU_CODEGEN_ARGS=--help` to see the full list of options)
+- [PR#940](https://github.com/EmbarkStudios/rust-gpu/pull/940) integrated the experimental [`SPIR-🇹` shader IR framework](https://github.com/EmbarkStudios/spirt) into the linker  
+  (opt-in via `RUSTGPU_CODEGEN_ARGS=--spirt`, see also [the `--spirt` docs](docs/src/codegen-args.md#--spirt), for more details)
 
 ### Changed 🛠️
 - [PR#958](https://github.com/EmbarkStudios/rust-gpu/pull/958) updated toolchain to `nightly-2022-10-29`

Cargo.lock

@@ -52,3 +52,7 @@ codegen-units = 256
 opt-level = 3
 incremental = true
 codegen-units = 256
+
+[patch.crates-io]
+# HACK(eddyb) only needed until `spirt 0.1.0` is released.
+spirt = { git = "https://github.com/EmbarkStudios/spirt.git", rev = "cda161b7e7b336685448ab0a7f8cfe96bd90e752" }

Cargo.toml

@@ -56,6 +56,7 @@ serde_json = "1.0"
 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"
 
 [dev-dependencies]
 pipe = "0.4"

crates/rustc_codegen_spirv/Cargo.toml

@@ -351,6 +351,12 @@ impl CodegenArgs {
             );
             opts.optflag("", "no-structurize", "disables CFG structurization");
 
+            opts.optflag(
+                "",
+                "spirt",
+                "use SPIR-T for legalization (see also `docs/src/codegen-args.md`)",
+            );
+
             // NOTE(eddyb) these are debugging options that used to be env vars
             // (for more information see `docs/src/codegen-args.md`).
             opts.optopt(
@@ -365,6 +371,12 @@ impl CodegenArgs {
                 "dump modules immediately after multimodule splitting, to files in DIR",
                 "DIR",
             );
+            opts.optopt(
+                "",
+                "dump-spirt-passes",
+                "dump the SPIR-T module across passes, to FILE and FILE.html",
+                "FILE",
+            );
             opts.optflag(
                 "",
                 "specializer-debug",
@@ -511,6 +523,7 @@ impl CodegenArgs {
             dce: !matches.opt_present("no-dce"),
             compact_ids: !matches.opt_present("no-compact-ids"),
             structurize: !matches.opt_present("no-structurize"),
+            spirt: matches.opt_present("spirt"),
 
             // FIXME(eddyb) deduplicate between `CodegenArgs` and `linker::Options`.
             emit_multiple_modules: module_output_type == ModuleOutputType::Multiple,
@@ -521,6 +534,7 @@ impl CodegenArgs {
             // (for more information see `docs/src/codegen-args.md`).
             dump_post_merge: matches_opt_path("dump-post-merge"),
             dump_post_split: matches_opt_dump_dir_path("dump-post-split"),
+            dump_spirt_passes: matches_opt_path("dump-spirt-passes"),
             specializer_debug: matches.opt_present("specializer-debug"),
             specializer_dump_instances: matches_opt_path("specializer-dump-instances"),
             print_all_zombie: matches.opt_present("print-all-zombie"),

crates/rustc_codegen_spirv/src/codegen_cx/mod.rs

@@ -34,6 +34,7 @@ pub struct Options {
     pub compact_ids: bool,
     pub dce: bool,
     pub structurize: bool,
+    pub spirt: bool,
 
     pub emit_multiple_modules: bool,
     pub spirv_metadata: SpirvMetadata,
@@ -48,6 +49,7 @@ pub struct Options {
     // (for more information see `docs/src/codegen-args.md`).
     pub dump_post_merge: Option<PathBuf>,
     pub dump_post_split: Option<PathBuf>,
+    pub dump_spirt_passes: Option<PathBuf>,
     pub specializer_debug: bool,
     pub specializer_dump_instances: Option<PathBuf>,
     pub print_all_zombie: bool,
@@ -141,10 +143,10 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
             bound += module.header.as_ref().unwrap().bound - 1;
             let this_version = module.header.as_ref().unwrap().version();
             if version != this_version {
-                sess.fatal(format!(
+                return Err(sess.err(format!(
                     "cannot link two modules with different SPIR-V versions: v{}.{} and v{}.{}",
                     version.0, version.1, this_version.0, this_version.1
-                ))
+                )));
             }
         }
 
@@ -234,25 +236,70 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
         );
     }
 
+    // NOTE(eddyb) with SPIR-T, we can do `mem2reg` before inlining, too!
+    if opts.spirt {
+        if opts.dce {
+            let _timer = sess.timer("link_dce-before-inlining");
+            dce::dce(&mut output);
+        }
+
+        let _timer = sess.timer("link_block_ordering_pass_and_mem2reg-before-inlining");
+        let mut pointer_to_pointee = FxHashMap::default();
+        let mut constants = FxHashMap::default();
+        let mut u32 = None;
+        for inst in &output.types_global_values {
+            match inst.class.opcode {
+                Op::TypePointer => {
+                    pointer_to_pointee
+                        .insert(inst.result_id.unwrap(), inst.operands[1].unwrap_id_ref());
+                }
+                Op::TypeInt
+                    if inst.operands[0].unwrap_literal_int32() == 32
+                        && inst.operands[1].unwrap_literal_int32() == 0 =>
+                {
+                    assert!(u32.is_none());
+                    u32 = Some(inst.result_id.unwrap());
+                }
+                Op::Constant if u32.is_some() && inst.result_type == u32 => {
+                    let value = inst.operands[0].unwrap_literal_int32();
+                    constants.insert(inst.result_id.unwrap(), value);
+                }
+                _ => {}
+            }
+        }
+        for func in &mut output.functions {
+            simple_passes::block_ordering_pass(func);
+            // Note: mem2reg requires functions to be in RPO order (i.e. block_ordering_pass)
+            mem2reg::mem2reg(
+                output.header.as_mut().unwrap(),
+                &mut output.types_global_values,
+                &pointer_to_pointee,
+                &constants,
+                func,
+            );
+            destructure_composites::destructure_composites(func);
+        }
+    }
+
     {
         let _timer = sess.timer("link_inline");
         inline::inline(sess, &mut output)?;
     }
 
     if opts.dce {
-        let _timer = sess.timer("link_dce");
+        let _timer = sess.timer("link_dce-after-inlining");
         dce::dce(&mut output);
     }
 
-    let mut output = if opts.structurize {
+    let mut output = if opts.structurize && !opts.spirt {
         let _timer = sess.timer("link_structurize");
         structurizer::structurize(output)
     } else {
         output
     };
 
     {
-        let _timer = sess.timer("link_block_ordering_pass_and_mem2reg");
+        let _timer = sess.timer("link_block_ordering_pass_and_mem2reg-after-inlining");
         let mut pointer_to_pointee = FxHashMap::default();
         let mut constants = FxHashMap::default();
         let mut u32 = None;
@@ -290,6 +337,86 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
         }
     }
 
+    if opts.spirt {
+        let mut per_pass_module_for_dumping = vec![];
+        let mut after_pass = |pass, module: &spirt::Module| {
+            if opts.dump_spirt_passes.is_some() {
+                per_pass_module_for_dumping.push((pass, module.clone()));
+            }
+        };
+
+        let spv_bytes = {
+            let _timer = sess.timer("assemble-to-spv_bytes-for-spirt");
+            spirv_tools::binary::from_binary(&output.assemble()).to_vec()
+        };
+        let cx = std::rc::Rc::new(spirt::Context::new());
+        let mut module = {
+            let _timer = sess.timer("spirt::Module::lower_from_spv_file");
+            match spirt::Module::lower_from_spv_bytes(cx.clone(), spv_bytes) {
+                Ok(module) => module,
+                Err(e) => {
+                    use rspirv::binary::Disassemble;
+
+                    return Err(sess
+                        .struct_err(format!("{e}"))
+                        .note(format!(
+                            "while lowering this SPIR-V module to SPIR-T:\n{}",
+                            output.disassemble()
+                        ))
+                        .emit());
+                }
+            }
+        };
+        after_pass("lower_from_spv", &module);
+
+        if opts.structurize {
+            {
+                let _timer = sess.timer("spirt::legalize::structurize_func_cfgs");
+                spirt::passes::legalize::structurize_func_cfgs(&mut module);
+            }
+            after_pass("structurize_func_cfgs", &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_file) = &opts.dump_spirt_passes {
+            // HACK(eddyb) using `.with_extension(...)` may replace part of a
+            // `.`-containing file name, but we want to always append `.html`.
+            let mut dump_file_html = dump_file.as_os_str().to_owned();
+            dump_file_html.push(".html");
+
+            let plan = spirt::print::Plan::for_versions(
+                &cx,
+                per_pass_module_for_dumping
+                    .iter()
+                    .map(|(pass, module)| (format!("after {pass}"), module)),
+            );
+            let pretty = plan.pretty_print();
+
+            // FIXME(eddyb) don't allocate whole `String`s here.
+            std::fs::write(dump_file, pretty.to_string()).unwrap();
+            std::fs::write(
+                dump_file_html,
+                pretty
+                    .render_to_html()
+                    .with_dark_mode_support()
+                    .to_html_doc(),
+            )
+            .unwrap();
+        }
+
+        let spv_words = {
+            let _timer = sess.timer("spirt::Module::lift_to_spv_module_emitter");
+            module.lift_to_spv_module_emitter().unwrap().words
+        };
+        output = {
+            let _timer = sess.timer("parse-spv_words-from-spirt");
+            let mut loader = Loader::new();
+            rspirv::binary::parse_words(&spv_words, &mut loader).unwrap();
+            loader.module()
+        };
+    }
+
     {
         let _timer = sess.timer("peephole_opts");
         let types = peephole_opts::collect_types(&output);

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -138,3 +138,18 @@ anyway, be careful).
 ### `--no-structurize`
 
 Disables CFG structurization. Probably results in invalid modules.
+
+### `--spirt`
+
+Enables using the experimental [`SPIR-🇹` shader IR framework](https://github.com/EmbarkStudios/spirt) in the linker - more specifically, this:
+- adds a `SPIR-V -> SPIR-🇹 -> SPIR-V` roundtrip  
+  (future `SPIR-🇹` passes would go in the middle of this, and eventually codegen might not produce `SPIR-V` at all)
+- replaces the existing structurizer with `SPIR-🇹` structurization (which is more robust and can e.g. handle `OpPhi`s)
+- runs some existing `SPIR-V` legalization/optimization passes (`mem2reg`) *before* inlining, instead of *only after* (as the `OpPhi`s they would produce are no longer an issue for structurization)
+
+For more information, also see [the `SPIR-🇹` repository](https://github.com/EmbarkStudios/spirt).
+
+### `--dump-spirt-passes FILE`
+
+Dump the `SPIR-🇹` module across passes (i.e. all of the versions before/after each pass), as a combined report, to `FILE` and `FILE.html`.  
+<sub>(the `.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>