Update spirt for DataInstForm interning.

Cargo.lock

@@ -462,6 +462,12 @@ dependencies = [
  "web-sys",
 ]
 
+[[package]]
+name = "convert_case"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6245d59a3e82a7fc217c5828a6692dbc6dfb63a0c8c90495621f7b9d79704a0e"
+
 [[package]]
 name = "core-foundation"
 version = "0.9.3"
@@ -594,6 +600,19 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "derive_more"
+version = "0.99.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4fb810d30a7c1953f91334de7244731fc3f3c10d7fe163338a35b9f640960321"
+dependencies = [
+ "convert_case",
+ "proc-macro2",
+ "quote",
+ "rustc_version",
+ "syn",
+]
+
 [[package]]
 name = "diff"
 version = "0.1.12"
@@ -2032,6 +2051,15 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "rustc_version"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bfa0f585226d2e68097d4f95d113b15b83a82e819ab25717ec0590d9584ef366"
+dependencies = [
+ "semver",
+]
+
 [[package]]
 name = "rustfix"
 version = "0.6.1"
@@ -2145,6 +2173,12 @@ dependencies = [
  "version-compare",
 ]
 
+[[package]]
+name = "semver"
+version = "1.0.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bebd363326d05ec3e2f532ab7660680f3b02130d780c299bca73469d521bc0ed"
+
 [[package]]
 name = "serde"
 version = "1.0.156"
@@ -2246,11 +2280,11 @@ dependencies = [
 [[package]]
 name = "spirt"
 version = "0.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
+source = "git+https://github.com/EmbarkStudios/spirt?branch=main#e8b4a474b2de791a27ea793a1d9c932c27a667ef"
-checksum = "e24fa996f12f3c667efbceaa99c222b8910a295a14d2c43c3880dfab2752def7"
 dependencies = [
  "arrayvec",
  "bytemuck",
+ "derive_more",
  "elsa",
  "indexmap",
  "internal-iterator",

Cargo.toml

@@ -52,3 +52,6 @@ codegen-units = 256
 opt-level = 3
 incremental = true
 codegen-units = 256
+
+[patch.crates-io]
+spirt = { git = "https://github.com/EmbarkStudios/spirt", branch = "main" }

crates/rustc_codegen_spirv/src/linker/spirt_passes/controlflow.rs

@@ -4,8 +4,8 @@ use crate::custom_insts::{self, CustomInst, CustomOp};
 use smallvec::SmallVec;
 use spirt::func_at::FuncAt;
 use spirt::{
-    cfg, spv, Attr, AttrSet, ConstCtor, ConstDef, ControlNodeKind, DataInstKind, DeclDef,
+    cfg, spv, Attr, AttrSet, ConstCtor, ConstDef, ControlNodeKind, DataInstFormDef, DataInstKind,
-    EntityDefs, ExportKey, Exportee, Module, Type, TypeCtor, TypeCtorArg, TypeDef, Value,
+    DeclDef, EntityDefs, ExportKey, Exportee, Module, Type, TypeCtor, TypeCtorArg, TypeDef, Value,
 };
 use std::fmt::Write as _;
 
@@ -108,14 +108,15 @@ pub fn convert_custom_aborts_to_unstructured_returns_in_entry_points(
                 .into_iter()
                 .filter_map(|func_at_inst| {
                     let data_inst_def = func_at_inst.def();
-                    if let DataInstKind::SpvInst(spv_inst) = &data_inst_def.kind {
+                    let data_inst_form_def = &cx[data_inst_def.form];
+                    if let DataInstKind::SpvInst(spv_inst) = &data_inst_form_def.kind {
                         if spv_inst.opcode == wk.OpLoad {
                             if let Value::Const(ct) = data_inst_def.inputs[0] {
                                 if let ConstCtor::PtrToGlobalVar(gv) = cx[ct].ctor {
                                     if interface_global_vars.contains(&gv) {
                                         return Some((
                                             gv,
-                                            data_inst_def.output_type.unwrap(),
+                                            data_inst_form_def.output_type.unwrap(),
                                             Value::DataInstOutput(func_at_inst.position),
                                         ));
                                     }
@@ -223,7 +224,7 @@ pub fn convert_custom_aborts_to_unstructured_returns_in_entry_points(
                 let data_inst_def = func_at_inst.def();
                 (
                     func_at_inst,
-                    match data_inst_def.kind {
+                    match cx[data_inst_def.form].kind {
                         DataInstKind::SpvExtInst { ext_set, inst }
                             if ext_set == custom_ext_inst_set =>
                         {
@@ -381,10 +382,13 @@ pub fn convert_custom_aborts_to_unstructured_returns_in_entry_points(
                         fmt += "\n";
 
                         let abort_inst_def = &mut func_def_body.data_insts[abort_inst];
-                        abort_inst_def.kind = DataInstKind::SpvExtInst {
+                        abort_inst_def.form = cx.intern(DataInstFormDef {
-                            ext_set: cx.intern("NonSemantic.DebugPrintf"),
+                            kind: DataInstKind::SpvExtInst {
-                            inst: 1,
+                                ext_set: cx.intern("NonSemantic.DebugPrintf"),
-                        };
+                                inst: 1,
+                            },
+                            output_type: cx[abort_inst_def.form].output_type,
+                        });
                         abort_inst_def.inputs = [Value::Const(mk_const_str(cx.intern(fmt)))]
                             .into_iter()
                             .chain(message_debug_printf_args)

crates/rustc_codegen_spirv/src/linker/spirt_passes/debuginfo.rs

@@ -25,6 +25,7 @@ pub fn convert_custom_debuginfo_to_spv(module: &mut Module) {
 
             seen_types: FxIndexSet::default(),
             seen_consts: FxIndexSet::default(),
+            seen_data_inst_forms: FxIndexSet::default(),
             seen_global_vars: FxIndexSet::default(),
             seen_funcs: FxIndexSet::default(),
         };
@@ -82,7 +83,7 @@ impl Transformer for CustomDebuginfoToSpv<'_> {
                 if let DataInstKind::SpvExtInst {
                     ext_set,
                     inst: ext_inst,
-                } = data_inst_def.kind
+                } = self.cx[data_inst_def.form].kind
                 {
                     if ext_set == self.custom_ext_inst_set {
                         let custom_op = CustomOp::decode(ext_inst);

crates/rustc_codegen_spirv/src/linker/spirt_passes/diagnostics.rs

@@ -11,8 +11,8 @@ use spirt::func_at::FuncAt;
 use spirt::visit::{InnerVisit, Visitor};
 use spirt::{
     spv, Attr, AttrSet, AttrSetDef, Const, ConstCtor, Context, ControlNode, ControlNodeKind,
-    DataInstDef, DataInstKind, Diag, DiagLevel, ExportKey, Exportee, Func, FuncDecl, GlobalVar,
+    DataInstDef, DataInstForm, DataInstKind, Diag, DiagLevel, ExportKey, Exportee, Func, FuncDecl,
-    InternedStr, Module, Type, Value,
+    GlobalVar, InternedStr, Module, Type, Value,
 };
 use std::marker::PhantomData;
 use std::{mem, str};
@@ -251,7 +251,7 @@ impl UseOrigin<'_> {
                     let wk = &super::SpvSpecWithExtras::get().well_known;
 
                     // FIXME(eddyb) deduplicate with `spirt_passes::diagnostics`.
-                    let custom_op = match debug_inst_def.kind {
+                    let custom_op = match cx[debug_inst_def.form].kind {
                         DataInstKind::SpvExtInst {
                             ext_set,
                             inst: ext_inst,
@@ -518,6 +518,11 @@ impl<'a> Visitor<'a> for DiagnosticReporter<'a> {
             }
         }
     }
+    fn visit_data_inst_form_use(&mut self, data_inst_form: DataInstForm) {
+        // NOTE(eddyb) this contains no deduplication because each `DataInstDef`
+        // will have any diagnostics reported separately.
+        self.visit_data_inst_form_def(&self.cx[data_inst_form]);
+    }
 
     fn visit_global_var_use(&mut self, gv: GlobalVar) {
         if self.seen_global_vars.insert(gv) {
@@ -617,7 +622,7 @@ impl<'a> Visitor<'a> for DiagnosticReporter<'a> {
                 if let DataInstKind::SpvExtInst {
                     ext_set,
                     inst: ext_inst,
-                } = data_inst_def.kind
+                } = self.cx[data_inst_def.form].kind
                 {
                     if ext_set == self.custom_ext_inst_set {
                         match CustomOp::decode(ext_inst) {
@@ -687,7 +692,7 @@ impl<'a> Visitor<'a> for DiagnosticReporter<'a> {
             _ => unreachable!(),
         }
 
-        if let DataInstKind::FuncCall(func) = data_inst_def.kind {
+        if let DataInstKind::FuncCall(func) = self.cx[data_inst_def.form].kind {
             // HACK(eddyb) visit `func` early, to control its `use_stack`, with
             // the later visit from `inner_visit_with` ignored as a duplicate.
             let old_origin = replace_origin(self, IntraFuncUseOrigin::CallCallee);

crates/rustc_codegen_spirv/src/linker/spirt_passes/mod.rs

@@ -13,8 +13,8 @@ use spirt::transform::InnerInPlaceTransform;
 use spirt::visit::{InnerVisit, Visitor};
 use spirt::{
     spv, AttrSet, Const, Context, ControlNode, ControlNodeKind, ControlRegion, DataInstDef,
-    DataInstKind, DeclDef, EntityOrientedDenseMap, Func, FuncDefBody, GlobalVar, Module, Type,
+    DataInstForm, DataInstFormDef, DataInstKind, DeclDef, EntityOrientedDenseMap, Func,
-    Value,
+    FuncDefBody, GlobalVar, Module, Type, Value,
 };
 use std::collections::VecDeque;
 use std::hash::Hash;
@@ -130,6 +130,7 @@ pub(super) fn run_func_passes<P>(
 
             seen_types: FxIndexSet::default(),
             seen_consts: FxIndexSet::default(),
+            seen_data_inst_forms: FxIndexSet::default(),
             seen_global_vars: FxIndexSet::default(),
             seen_funcs: FxIndexSet::default(),
         };
@@ -181,7 +182,7 @@ pub(super) fn run_func_passes<P>(
                 pass_fn(cx, func_def_body);
 
                 // FIXME(eddyb) avoid doing this except where changes occurred.
-                remove_unused_values_in_func(func_def_body);
+                remove_unused_values_in_func(cx, func_def_body);
             }
         }
         after_pass(full_name, module, profiler);
@@ -196,6 +197,7 @@ struct ReachableUseCollector<'a> {
     // FIXME(eddyb) build some automation to avoid ever repeating these.
     seen_types: FxIndexSet<Type>,
     seen_consts: FxIndexSet<Const>,
+    seen_data_inst_forms: FxIndexSet<DataInstForm>,
     seen_global_vars: FxIndexSet<GlobalVar>,
     seen_funcs: FxIndexSet<Func>,
 }
@@ -213,6 +215,11 @@ impl Visitor<'_> for ReachableUseCollector<'_> {
             self.visit_const_def(&self.cx[ct]);
         }
     }
+    fn visit_data_inst_form_use(&mut self, data_inst_form: DataInstForm) {
+        if self.seen_data_inst_forms.insert(data_inst_form) {
+            self.visit_data_inst_form_def(&self.cx[data_inst_form]);
+        }
+    }
 
     fn visit_global_var_use(&mut self, gv: GlobalVar) {
         if self.seen_global_vars.insert(gv) {
@@ -247,6 +254,7 @@ const _: () = {
         fn visit_attr_set_use(&mut self, _: AttrSet) {}
         fn visit_type_use(&mut self, _: Type) {}
         fn visit_const_use(&mut self, _: Const) {}
+        fn visit_data_inst_form_use(&mut self, _: DataInstForm) {}
         fn visit_global_var_use(&mut self, _: GlobalVar) {}
         fn visit_func_use(&mut self, _: Func) {}
 
@@ -354,7 +362,7 @@ const _: () = {
 /// a function body (both `DataInst`s and `ControlRegion` inputs/outputs).
 //
 // FIXME(eddyb) should this be a dedicated pass?
-fn remove_unused_values_in_func(func_def_body: &mut FuncDefBody) {
+fn remove_unused_values_in_func(cx: &Context, func_def_body: &mut FuncDefBody) {
     // Avoid having to support unstructured control-flow.
     if func_def_body.unstructured_cfg.is_some() {
         return;
@@ -473,7 +481,9 @@ fn remove_unused_values_in_func(func_def_body: &mut FuncDefBody) {
                             for func_at_inst in func_at_control_node.at(insts) {
                                 // Ignore pure instructions (i.e. they're only used
                                 // if their output value is used, from somewhere else).
-                                if let DataInstKind::SpvInst(spv_inst) = &func_at_inst.def().kind {
+                                if let DataInstKind::SpvInst(spv_inst) =
+                                    &cx[func_at_inst.def().form].kind
+                                {
                                     // HACK(eddyb) small selection relevant for now,
                                     // but should be extended using e.g. a bitset.
                                     if [wk.OpNop, wk.OpCompositeInsert].contains(&spv_inst.opcode) {
@@ -518,7 +528,9 @@ fn remove_unused_values_in_func(func_def_body: &mut FuncDefBody) {
                 let mut all_nops = true;
                 let mut func_at_inst_iter = func_def_body.at_mut(insts).into_iter();
                 while let Some(mut func_at_inst) = func_at_inst_iter.next() {
-                    if let DataInstKind::SpvInst(spv_inst) = &func_at_inst.reborrow().def().kind {
+                    if let DataInstKind::SpvInst(spv_inst) =
+                        &cx[func_at_inst.reborrow().def().form].kind
+                    {
                         if spv_inst.opcode == wk.OpNop {
                             continue;
                         }
@@ -528,10 +540,14 @@ fn remove_unused_values_in_func(func_def_body: &mut FuncDefBody) {
                     {
                         // Replace the removed `DataInstDef` itself with `OpNop`,
                         // removing the ability to use its "name" as a value.
+                        //
+                        // FIXME(eddyb) cache the interned `OpNop`.
                         *func_at_inst.def() = DataInstDef {
                             attrs: Default::default(),
-                            kind: DataInstKind::SpvInst(wk.OpNop.into()),
+                            form: cx.intern(DataInstFormDef {
-                            output_type: None,
+                                kind: DataInstKind::SpvInst(wk.OpNop.into()),
+                                output_type: None,
+                            }),
                             inputs: iter::empty().collect(),
                         };
                         continue;

crates/rustc_codegen_spirv/src/linker/spirt_passes/reduce.rs

@@ -6,8 +6,8 @@ use spirt::visit::InnerVisit;
 use spirt::{
     spv, Const, ConstCtor, ConstDef, Context, ControlNode, ControlNodeDef, ControlNodeKind,
     ControlNodeOutputDecl, ControlRegion, ControlRegionInputDecl, DataInst, DataInstDef,
-    DataInstKind, EntityOrientedDenseMap, FuncDefBody, SelectionKind, Type, TypeCtor, TypeDef,
+    DataInstFormDef, DataInstKind, EntityOrientedDenseMap, FuncDefBody, SelectionKind, Type,
-    Value,
+    TypeCtor, TypeDef, Value,
 };
 use std::collections::hash_map::Entry;
 use std::convert::{TryFrom, TryInto};
@@ -65,7 +65,7 @@ pub(crate) fn reduce_in_func(cx: &Context, func_def_body: &mut FuncDefBody) {
                     ..
                 } => {
                     for func_at_inst in func_at_control_node.at(insts) {
-                        if let Ok(redu) = Reducible::try_from(func_at_inst.def()) {
+                        if let Ok(redu) = Reducible::try_from((cx, func_at_inst.def())) {
                             let redu_target = ReductionTarget::DataInst(func_at_inst.position);
                             reduction_queue.push((redu_target, redu));
                         }
@@ -202,10 +202,14 @@ pub(crate) fn reduce_in_func(cx: &Context, func_def_body: &mut FuncDefBody) {
 
                         // Replace the reduced `DataInstDef` itself with `OpNop`,
                         // removing the ability to use its "name" as a value.
+                        //
+                        // FIXME(eddyb) cache the interned `OpNop`.
                         *func_def_body.at_mut(inst).def() = DataInstDef {
                             attrs: Default::default(),
-                            kind: DataInstKind::SpvInst(wk.OpNop.into()),
+                            form: cx.intern(DataInstFormDef {
-                            output_type: None,
+                                kind: DataInstKind::SpvInst(wk.OpNop.into()),
+                                output_type: None,
+                            }),
                             inputs: iter::empty().collect(),
                         };
                     }
@@ -499,12 +503,14 @@ impl<V> Reducible<V> {
     }
 }
 
-impl TryFrom<&DataInstDef> for Reducible {
+// FIXME(eddyb) instead of taking a `&Context`, could `Reducible` hold a `DataInstForm`?
+impl TryFrom<(&Context, &DataInstDef)> for Reducible {
     type Error = ();
-    fn try_from(inst_def: &DataInstDef) -> Result<Self, ()> {
+    fn try_from((cx, inst_def): (&Context, &DataInstDef)) -> Result<Self, ()> {
-        if let DataInstKind::SpvInst(spv_inst) = &inst_def.kind {
+        let inst_form_def = &cx[inst_def.form];
+        if let DataInstKind::SpvInst(spv_inst) = &inst_form_def.kind {
             let op = PureOp::try_from(spv_inst)?;
-            let output_type = inst_def.output_type.unwrap();
+            let output_type = inst_form_def.output_type.unwrap();
             if let [input] = inst_def.inputs[..] {
                 return Ok(Self {
                     op,
@@ -517,15 +523,21 @@ impl TryFrom<&DataInstDef> for Reducible {
     }
 }
 
-// HACK(eddyb) `IntToBool` is the only reason this is `TryFrom` not `From`.
+impl Reducible {
-impl TryFrom<Reducible> for DataInstDef {
+    // HACK(eddyb) `IntToBool` is the only reason this can return `None`.
-    type Error = ();
+    fn try_into_inst(self, cx: &Context) -> Option<DataInstDef> {
-    fn try_from(redu: Reducible) -> Result<Self, ()> {
+        let Self {
-        Ok(Self {
+            op,
+            output_type,
+            input,
+        } = self;
+        Some(DataInstDef {
             attrs: Default::default(),
-            kind: DataInstKind::SpvInst(redu.op.try_into()?),
+            form: cx.intern(DataInstFormDef {
-            output_type: Some(redu.output_type),
+                kind: DataInstKind::SpvInst(op.try_into().ok()?),
-            inputs: iter::once(redu.input).collect(),
+                output_type: Some(output_type),
+            }),
+            inputs: iter::once(input).collect(),
         })
     }
 }
@@ -624,11 +636,11 @@ enum ReductionStep {
 
 impl Reducible<&DataInstDef> {
     // FIXME(eddyb) force the input to actually be itself some kind of pure op.
-    fn try_reduce_output_of_data_inst(&self) -> Option<ReductionStep> {
+    fn try_reduce_output_of_data_inst(&self, cx: &Context) -> Option<ReductionStep> {
         let wk = &super::SpvSpecWithExtras::get().well_known;
 
         let input_inst_def = self.input;
-        if let DataInstKind::SpvInst(input_spv_inst) = &input_inst_def.kind {
+        if let DataInstKind::SpvInst(input_spv_inst) = &cx[input_inst_def.form].kind {
             // NOTE(eddyb) do not destroy information left in e.g. comments.
             #[allow(clippy::match_same_arms)]
             match self.op {
@@ -737,8 +749,9 @@ impl Reducible {
                     .try_reduce(cx, func.reborrow(), value_replacements, parent_map, cache)?;
                 // HACK(eddyb) this is here because it can fail, see the comment
                 // on `output_from_updated_state` for what's actually going on.
-                let output_from_updated_state_inst =
+                let output_from_updated_state_inst = self
-                    DataInstDef::try_from(self.with_input(input_from_updated_state)).ok()?;
+                    .with_input(input_from_updated_state)
+                    .try_into_inst(cx)?;
 
                 // Now that the reduction succeeded for the initial state,
                 // we can proceed with augmenting the loop with the extra state.
@@ -874,7 +887,10 @@ impl Reducible {
             }
             Value::DataInstOutput(inst) => {
                 let inst_def = &*func.reborrow().at(inst).def();
-                match self.with_input(inst_def).try_reduce_output_of_data_inst()? {
+                match self
+                    .with_input(inst_def)
+                    .try_reduce_output_of_data_inst(cx)?
+                {
                     ReductionStep::Complete(v) => Some(v),
                     // FIXME(eddyb) actually use a loop instead of recursing here.
                     ReductionStep::Partial(redu) => {