linker/duplicates: handle all decorations instead of special-casing zombies.

This commit is contained in:
Eduard-Mihai Burtescu 2022-11-30 03:46:39 +02:00 committed by Eduard-Mihai Burtescu
parent f78b23b926
commit a9ede2ed9f
2 changed files with 52 additions and 28 deletions

View File

@ -1,4 +1,3 @@
use crate::decorations::{CustomDecoration, ZombieDecoration};
use rspirv::binary::Assemble;
use rspirv::dr::{Instruction, Module, Operand};
use rspirv::spirv::{Op, Word};
@ -73,15 +72,20 @@ fn make_annotation_key(inst: &Instruction) -> Vec<u32> {
fn gather_annotations(annotations: &[Instruction]) -> FxHashMap<Word, Vec<u32>> {
let mut map = FxHashMap::default();
for inst in annotations {
if inst.class.opcode == Op::Decorate || inst.class.opcode == Op::MemberDecorate {
match map.entry(inst.operands[0].id_ref_any().unwrap()) {
match inst.class.opcode {
Op::Decorate
| Op::DecorateId
| Op::DecorateString
| Op::MemberDecorate
| Op::MemberDecorateString => match map.entry(inst.operands[0].id_ref_any().unwrap()) {
hash_map::Entry::Vacant(entry) => {
entry.insert(vec![make_annotation_key(inst)]);
}
hash_map::Entry::Occupied(mut entry) => {
entry.get_mut().push(make_annotation_key(inst));
}
}
},
_ => {}
}
}
map.into_iter()
@ -110,7 +114,6 @@ fn gather_names(debug_names: &[Instruction]) -> FxHashMap<Word, String> {
fn make_dedupe_key(
inst: &Instruction,
unresolved_forward_pointers: &FxHashSet<Word>,
zombies: &FxHashSet<Word>,
annotations: &FxHashMap<Word, Vec<u32>>,
names: &FxHashMap<Word, String>,
) -> Vec<u32> {
@ -136,15 +139,6 @@ fn make_dedupe_key(
}
}
if let Some(id) = inst.result_id {
data.push(if zombies.contains(&id) {
if inst.result_type.is_some() {
id
} else {
1
}
} else {
0
});
if let Some(annos) = annotations.get(&id) {
data.extend_from_slice(annos);
}
@ -152,6 +146,9 @@ fn make_dedupe_key(
// Names only matter for OpVariable.
if let Some(name) = names.get(&id) {
// Jump through some hoops to shove a String into a Vec<u32>.
//
// FIXME(eddyb) this should `.assemble_into(&mut data)` the
// `Operand::LiteralString(...)` from the original `Op::Name`.
for chunk in name.as_bytes().chunks(4) {
let slice = match *chunk {
[a] => [a, 0, 0, 0],
@ -184,11 +181,6 @@ fn rewrite_inst_with_rules(inst: &mut Instruction, rules: &FxHashMap<u32, u32>)
pub fn remove_duplicate_types(module: &mut Module) {
// Keep in mind, this algorithm requires forward type references to not exist - i.e. it's a valid spir-v module.
// Include zombies in the key to not merge zombies with non-zombies
let zombies: FxHashSet<Word> = ZombieDecoration::decode_all(module)
.map(|(z, _)| z)
.collect();
// When a duplicate type is encountered, then this is a map from the deleted ID, to the new, deduplicated ID.
let mut rewrite_rules = FxHashMap::default();
// Instructions are encoded into "keys": their opcode, followed by arguments, then annotations.
@ -224,13 +216,7 @@ pub fn remove_duplicate_types(module: &mut Module) {
// all_inst_iter_mut pass below. However, the code is a lil bit cleaner this way I guess.
rewrite_inst_with_rules(inst, &rewrite_rules);
let key = make_dedupe_key(
inst,
&unresolved_forward_pointers,
&zombies,
&annotations,
&names,
);
let key = make_dedupe_key(inst, &unresolved_forward_pointers, &annotations, &names);
match key_to_result_id.entry(key) {
hash_map::Entry::Vacant(entry) => {

View File

@ -1,4 +1,5 @@
use super::Result;
use crate::decorations::{CustomDecoration, ZombieDecoration};
use rspirv::dr::{Instruction, Module};
use rspirv::spirv::{Capability, Decoration, LinkageType, Op, Word};
use rustc_data_structures::fx::{FxHashMap, FxHashSet};
@ -26,6 +27,27 @@ fn find_import_export_pairs_and_killed_params(
let mut rewrite_rules = FxHashMap::default();
let mut killed_parameters = FxHashSet::default();
// HACK(eddyb) collect all zombies, and anything that transitively mentions
// them (or is "infected"), to ignore them when doing type checks, as they
// will either be removed later, or become an error of their own.
let mut zombie_infected: FxHashSet<Word> = ZombieDecoration::decode_all(module)
.map(|(z, _)| z)
.collect();
for inst in module.global_inst_iter() {
if let Some(result_id) = inst.result_id {
if !zombie_infected.contains(&result_id) {
let mut id_operands = inst.operands.iter().filter_map(|o| o.id_ref_any());
// NOTE(eddyb) this takes advantage of the fact that the module
// is ordered def-before-use (with the minor exception of forward
// references for recursive data, which are not fully supported),
// to be able to propagate "zombie infection" in one pass.
if id_operands.any(|id| zombie_infected.contains(&id)) {
zombie_infected.insert(result_id);
}
}
}
}
// First, collect all the exports.
for annotation in &module.annotations {
let (id, name) = match get_linkage_inst(annotation) {
@ -53,7 +75,14 @@ fn find_import_export_pairs_and_killed_params(
};
let import_type = *type_map.get(&import_id).expect("Unexpected op");
// Make sure the import/export pair has the same type.
check_tys_equal(sess, module, name, import_type, export_type)?;
check_tys_equal(
sess,
module,
name,
import_type,
export_type,
&zombie_infected,
)?;
rewrite_rules.insert(import_id, export_id);
if let Some(params) = fn_parameters.get(&import_id) {
for &param in params {
@ -111,8 +140,17 @@ fn check_tys_equal(
name: &str,
import_type: Word,
export_type: Word,
zombie_infected: &FxHashSet<Word>,
) -> Result<()> {
if import_type == export_type {
let allowed = import_type == export_type || {
// HACK(eddyb) zombies can cause types to differ in definition, due to
// requiring multiple different instances (usually different `Span`s),
// so we ignore them, as `find_import_export_pairs_and_killed_params`'s
// own comment explains (zombies will cause errors or be removed, *later*).
zombie_infected.contains(&import_type) && zombie_infected.contains(&export_type)
};
if allowed {
Ok(())
} else {
// We have an error. It's okay to do something really slow now to report the error.