spirt-passes: add reduce pass for replacing ops with their inputs/c…

…onstants.
EmbarkStudios · Jan 16, 2023 · 235474b · 235474b
1 parent 309739d
commit 235474b
Show file tree

Hide file tree

Showing 4 changed files with 931 additions and 1 deletion.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/rustc_codegen_spirv/Cargo.toml b/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"

diff --git a/crates/rustc_codegen_spirv/src/linker/spirt_passes/mod.rs b/crates/rustc_codegen_spirv/src/linker/spirt_passes/mod.rs
@@ -1,8 +1,83 @@
 //! SPIR-T pass infrastructure and supporting utilities.
 
+mod reduce;
+
+use lazy_static::lazy_static;
 use rustc_data_structures::fx::FxIndexSet;
 use spirt::visit::{InnerVisit, Visitor};
-use spirt::{AttrSet, Const, Context, DeclDef, Func, GlobalVar, Module, Type};
+use spirt::{spv, AttrSet, Const, Context, DeclDef, Func, GlobalVar, Module, Type};
+
+// HACK(eddyb) `spv::spec::Spec` with extra `WellKnown`s (that should be upstreamed).
+macro_rules! def_spv_spec_with_extra_well_known {
+    ($($group:ident: $ty:ty = [$($entry:ident),+ $(,)?]),+ $(,)?) => {
+        struct SpvSpecWithExtras {
+            __base_spec: &'static spv::spec::Spec,
+
+            well_known: SpvWellKnownWithExtras,
+        }
+
+        #[allow(non_snake_case)]
+        pub struct SpvWellKnownWithExtras {
+            __base_well_known: &'static spv::spec::WellKnown,
+
+            $($(pub $entry: $ty,)+)+
+        }
+
+        impl std::ops::Deref for SpvSpecWithExtras {
+            type Target = spv::spec::Spec;
+            fn deref(&self) -> &Self::Target {
+                self.__base_spec
+            }
+        }
+
+        impl std::ops::Deref for SpvWellKnownWithExtras {
+            type Target = spv::spec::WellKnown;
+            fn deref(&self) -> &Self::Target {
+                self.__base_well_known
+            }
+        }
+
+        impl SpvSpecWithExtras {
+            #[inline(always)]
+            #[must_use]
+            pub fn get() -> &'static SpvSpecWithExtras {
+                lazy_static! {
+                    static ref SPEC: SpvSpecWithExtras = {
+                        #[allow(non_camel_case_types)]
+                        struct PerWellKnownGroup<$($group),+> {
+                            $($group: $group),+
+                        }
+
+                        let spv_spec = spv::spec::Spec::get();
+                        let lookup_fns = PerWellKnownGroup {
+                            opcode: |name| spv_spec.instructions.lookup(name).unwrap(),
+                        };
+
+                        SpvSpecWithExtras {
+                            __base_spec: spv_spec,
+
+                            well_known: SpvWellKnownWithExtras {
+                                __base_well_known: &spv_spec.well_known,
+
+                                $($($entry: (lookup_fns.$group)(stringify!($entry)),)+)+
+                            },
+                        }
+                    };
+                }
+                &SPEC
+            }
+        }
+    };
+}
+def_spv_spec_with_extra_well_known! {
+    opcode: spv::spec::Opcode = [
+        OpConstantComposite,
+
+        OpBitcast,
+        OpCompositeInsert,
+        OpCompositeExtract,
+    ],
+}
 
 /// Run intra-function passes on all `Func` definitions in the `Module`.
 //
@@ -36,6 +111,7 @@ pub(super) fn run_func_passes<P>(
     for name in passes {
         let name = name.as_ref();
         let (full_name, pass_fn) = match name {
+            "reduce" => ("spirt_passes::reduce", reduce::reduce_in_func),
             _ => panic!("unknown `--spirt-passes={}`", name),
         };