From 0557036fb109bc3c05bff80753f0efc21ecb448f Mon Sep 17 00:00:00 2001 From: Eduard-Mihai Burtescu Date: Thu, 1 Jun 2023 18:34:10 +0300 Subject: [PATCH] Preserve all intra-function annotations in both `specializer` and `inline`. --- .../rustc_codegen_spirv/src/linker/inline.rs | 19 +++++++++---------- .../src/linker/specializer.rs | 16 +++++++++++++++- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/crates/rustc_codegen_spirv/src/linker/inline.rs b/crates/rustc_codegen_spirv/src/linker/inline.rs index 223137b542..144a1e98d2 100644 --- a/crates/rustc_codegen_spirv/src/linker/inline.rs +++ b/crates/rustc_codegen_spirv/src/linker/inline.rs @@ -354,17 +354,16 @@ impl Inliner<'_, '_> { /// Applies all rewrite rules to the decorations in the header. fn apply_rewrite_for_decorations(&mut self, rewrite_rules: &FxHashMap) { - // NOTE(siebencorgie): We don't care *what* decoration we rewrite atm. AFAIK there is no case where rewriting - // the decoration on inline wouldn't be valid. + // NOTE(siebencorgie): We don't care *what* decoration we rewrite atm. + // AFAIK there is no case where keeping decorations on inline wouldn't be valid. for annotation_idx in 0..self.annotations.len() { - if self.annotations[annotation_idx].class.opcode == Op::Decorate { - if let Some(id) = self.annotations[annotation_idx].operands[0].id_ref_any_mut() { - if let Some(&rewrite) = rewrite_rules.get(id) { - // Copy decoration instruction and push it. - let mut instcpy = self.annotations[annotation_idx].clone(); - *instcpy.operands[0].id_ref_any_mut().unwrap() = rewrite; - self.annotations.push(instcpy); - } + let inst = &self.annotations[annotation_idx]; + if let [Operand::IdRef(target), ..] = inst.operands[..] { + if let Some(&rewritten_target) = rewrite_rules.get(&target) { + // Copy decoration instruction and push it. + let mut cloned_inst = inst.clone(); + cloned_inst.operands[0] = Operand::IdRef(rewritten_target); + self.annotations.push(cloned_inst); } } } diff --git a/crates/rustc_codegen_spirv/src/linker/specializer.rs b/crates/rustc_codegen_spirv/src/linker/specializer.rs index 582daf1da9..d1cf59ba39 100644 --- a/crates/rustc_codegen_spirv/src/linker/specializer.rs +++ b/crates/rustc_codegen_spirv/src/linker/specializer.rs @@ -2379,7 +2379,7 @@ impl<'a, S: Specialization> Expander<'a, S> { let expanded_debug_names = expand_debug_or_annotation(debug_names); // Expand `Op(Member)Decorate* %target ...`, when `target` is "generic". - let expanded_annotations = expand_debug_or_annotation(annotations); + let mut expanded_annotations = expand_debug_or_annotation(annotations); // Expand "generic" globals (types, constants and module-scoped variables). let mut expanded_types_global_values = @@ -2440,6 +2440,8 @@ impl<'a, S: Specialization> Expander<'a, S> { let mut rewrite_rules = FxHashMap::default(); for func in newly_expanded_functions { + rewrite_rules.clear(); + rewrite_rules.extend(func.parameters.iter_mut().map(|param| { let old_id = param.result_id.unwrap(); let new_id = self.builder.id(); @@ -2460,6 +2462,18 @@ impl<'a, S: Specialization> Expander<'a, S> { ); super::apply_rewrite_rules(&rewrite_rules, &mut func.blocks); + + // HACK(eddyb) this duplicates similar logic from `inline`. + for annotation_idx in 0..expanded_annotations.len() { + let inst = &expanded_annotations[annotation_idx]; + if let [Operand::IdRef(target), ..] = inst.operands[..] { + if let Some(&rewritten_target) = rewrite_rules.get(&target) { + let mut expanded_inst = inst.clone(); + expanded_inst.operands[0] = Operand::IdRef(rewritten_target); + expanded_annotations.push(expanded_inst); + } + } + } } }