Linker perf (#33)

* Improve linker performance - this should loop over data significantly less

* Remove Rc and RefCell because they weren't needed in the end

* Moved some code into a DefUseAnalyer

* Point to same rspirv

* Link to rspirv git instead of locally

Cargo.toml

@@ -4,3 +4,6 @@ members = [
     "rustc_codegen_spirv",
     "spirv-std",
 ]
+
+[patch.crates-io]
+rspirv = { git = "https://github.com/gfx-rs/rspirv.git", rev = "1addc7d33ae1460ffa683e2e6311e466ac876c23" }

rspirv-linker/Cargo.toml

@@ -7,7 +7,7 @@ edition = "2018"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-rspirv = { git = "https://github.com/gfx-rs/rspirv" }
+rspirv = "0.7.0"
 topological-sort = "0.1"
 thiserror = "1.0.20"
 

rspirv-linker/src/lib.rs

@@ -167,47 +167,119 @@ fn kill_annotations_and_debug(module: &mut rspirv::dr::Module, id: u32) {
     kill_with_id(&mut module.debugs, id);
 }
 
-fn remove_duplicate_types(module: &mut rspirv::dr::Module) {
+fn remove_duplicate_types(module: rspirv::dr::Module) -> rspirv::dr::Module {
+    use rspirv::binary::Assemble;
+
     // jb-todo: spirv-tools's linker has special case handling for SpvOpTypeForwardPointer,
     // not sure if we need that; see https://github.com/KhronosGroup/SPIRV-Tools/blob/e7866de4b1dc2a7e8672867caeb0bdca49f458d3/source/opt/remove_duplicates_pass.cpp for reference
 
+    let mut instructions = module
+        .all_inst_iter()
+        .cloned()
+        .collect::<Vec<_>>()
+        .into_boxed_slice(); // force boxed slice so we don't accidentally grow or shrink it later
+
+    let mut def_use_analyzer = DefUseAnalyzer::new(&mut instructions);
+
+    let mut kill_annotations = vec![];
+    let mut continue_from_idx = 0;
+
     // need to do this process iteratively because types can reference each other
     loop {
-        let mut replace = None;
-
-        // start with `nth` so we can restart this loop quickly after killing the op
-        for (i_idx, i) in module.types_global_values.iter().enumerate() {
-            let mut identical = None;
-            for j in module.types_global_values.iter().skip(i_idx + 1) {
-                if i.is_type_identical(j) {
-                    identical = j.result_id;
-                    break;
-                }
+        let mut dedup = std::collections::HashMap::new();
+        let mut duplicate = None;
+
+        for (iterator_idx, module_inst) in module
+            .types_global_values
+            .iter()
+            .enumerate()
+            .skip(continue_from_idx)
+        {
+            let (inst_idx, inst) = def_use_analyzer.def(module_inst.result_id.unwrap());
+
+            if inst.class.opcode == spirv::Op::Nop {
+                continue;
             }
 
-            if let Some(identical) = identical {
-                replace = Some((i.result_id.unwrap(), identical, i_idx));
+            // partially assemble only the opcode and operands to be used as a key
+            // maybe this should also include the result_type
+            let data = {
+                let mut data = vec![];
+
+                data.push(inst.class.opcode as u32);
+                for op in &inst.operands {
+                    op.assemble_into(&mut data);
+                }
+
+                data
+            };
+
+            // dedup contains a tuple of three indices;
+            // the first two point into our `def_use_analyzer.instructions` map
+            // the last one points into the `module.types_global_values` iterator so we can resume iteration
+            dedup
+                .entry(data)
+                .and_modify(|(identical_idx, backtrack_idx)| {
+                    duplicate = Some((inst_idx, *identical_idx, *backtrack_idx));
+                })
+                .or_insert((inst_idx, iterator_idx)); // store the index that we encountered an instruction
+                                                      // for the first time so we can backtrack later
+
+            if let Some((_, _, backtrack_idx)) = duplicate {
+                continue_from_idx = backtrack_idx;
                 break;
             }
         }
 
-        // can't do this directly in the previous loop because of the
-        // mut borrow needed on `module`
-        if let Some((remove, keep, kill_idx)) = replace {
-            kill_annotations_and_debug(module, remove);
-            replace_all_uses_with(module, remove, keep);
-            module.types_global_values.swap_remove(kill_idx);
+        if let Some((before_idx, after_idx, _)) = duplicate {
+            let before_id = def_use_analyzer.instructions[before_idx].result_id.unwrap();
+            let after_id = def_use_analyzer.instructions[after_idx].result_id.unwrap();
+
+            // remove annotations later
+            kill_annotations.push(before_id);
+
+            def_use_analyzer.for_each_use(before_id, |inst| {
+                if inst.result_type == Some(before_id) {
+                    inst.result_type = Some(after_id);
+                }
+
+                for op in inst.operands.iter_mut() {
+                    match op {
+                        rspirv::dr::Operand::IdMemorySemantics(w)
+                        | rspirv::dr::Operand::IdScope(w)
+                        | rspirv::dr::Operand::IdRef(w) => {
+                            if *w == before_id {
+                                *w = after_id
+                            }
+                        }
+                        _ => {}
+                    }
+                }
+            });
+
+            // this loop / system works on the assumption that all indices remain valid,
+            // so instead of removing the instruction we just nop it out - `consume_instruction` will then
+            // skip all OpNops and they won't appear in the newly constructed module
+            def_use_analyzer.instructions[before_idx] =
+                rspirv::dr::Instruction::new(spirv::Op::Nop, None, None, vec![]);
         } else {
             break;
         }
     }
-}
 
-fn remove_duplicates(module: &mut rspirv::dr::Module) {
-    remove_duplicate_capablities(module);
-    remove_duplicate_ext_inst_imports(module);
-    remove_duplicate_types(module);
-    // jb-todo: strip identical OpDecoration / OpDecorationGroups
+    let mut loader = rspirv::dr::Loader::new();
+
+    for inst in def_use_analyzer.instructions.iter() {
+        loader.consume_instruction(inst.clone());
+    }
+
+    let mut module = loader.module();
+
+    for remove in kill_annotations {
+        kill_annotations_and_debug(&mut module, remove);
+    }
+
+    module
 }
 
 #[derive(Clone, Debug)]
@@ -402,6 +474,89 @@ impl DefAnalyzer {
     }
 }
 
+struct DefUseAnalyzer<'a> {
+    def_ids: HashMap<u32, usize>,
+    use_ids: HashMap<u32, Vec<usize>>,
+    use_result_type_ids: HashMap<u32, Vec<usize>>,
+    instructions: &'a mut [rspirv::dr::Instruction]
+}
+
+impl<'a> DefUseAnalyzer<'a> {
+    fn new(instructions: &'a mut [rspirv::dr::Instruction]) -> Self{
+        let mut def_ids = HashMap::new();
+        let mut use_ids: HashMap<u32, Vec<usize>> = HashMap::new();
+        let mut use_result_type_ids: HashMap<u32, Vec<usize>> = HashMap::new();
+
+        instructions
+            .iter()
+            .enumerate()
+            .for_each(|(inst_idx, inst)| {
+                if let Some(def_id) = inst.result_id {
+                    def_ids
+                        .entry(def_id)
+                        .and_modify(|stored_inst| {
+                            *stored_inst = inst_idx;
+                        })
+                        .or_insert(inst_idx);
+                }
+
+                if let Some(result_type) = inst.result_type {
+                    use_result_type_ids
+                        .entry(result_type)
+                        .and_modify(|v| v.push(inst_idx))
+                        .or_insert(vec![inst_idx]);
+                }
+
+                for op in inst.operands.iter() {
+                    match op {
+                        rspirv::dr::Operand::IdMemorySemantics(w)
+                        | rspirv::dr::Operand::IdScope(w)
+                        | rspirv::dr::Operand::IdRef(w) => {
+                            use_ids
+                                .entry(*w)
+                                .and_modify(|v| v.push(inst_idx))
+                                .or_insert(vec![inst_idx]);
+                        }
+                        _ => {}
+                    }
+                }
+            });
+
+        Self {
+            def_ids,
+            use_ids,
+            use_result_type_ids,
+            instructions
+        }
+    }
+
+    fn def_idx(&self, id: u32) -> usize {
+        self.def_ids[&id]
+    }
+
+    fn def(&self, id: u32) -> (usize, &rspirv::dr::Instruction) {
+        let idx = self.def_idx(id);
+        (idx, &self.instructions[idx])
+    }
+
+    fn for_each_use<F>(&mut self, id: u32, mut f: F) 
+    where F: FnMut(&mut rspirv::dr::Instruction) {
+        // find by `result_type`
+        if let Some(use_result_type_id) = self.use_result_type_ids.get(&id) {
+            for inst_idx in use_result_type_id {
+                f(&mut self.instructions[*inst_idx])
+            }
+        }
+
+        // find by operand
+        if let Some(use_id) = self.use_ids.get(&id) {
+            for inst_idx in use_id {
+                f(&mut self.instructions[*inst_idx]);
+            }
+        }
+    }
+}
+
 fn import_kill_annotations_and_debug(module: &mut rspirv::dr::Module, info: &LinkInfo) {
     for import in &info.imports {
         kill_annotations_and_debug(module, import.id);
@@ -779,10 +934,6 @@ pub fn link(inputs: &mut [&mut rspirv::dr::Module], opts: &Options) -> Result<rs
         bound += module.header.as_ref().unwrap().bound - 1;
     }
 
-    for i in inputs.iter() {
-        println!("{}\n\n", i.disassemble());
-    }
-
     // merge the binaries
     let mut loader = rspirv::dr::Loader::new();
 
@@ -793,7 +944,6 @@ pub fn link(inputs: &mut [&mut rspirv::dr::Module], opts: &Options) -> Result<rs
     }
 
     let mut output = loader.module();
-    println!("{}\n\n", output.disassemble());
 
     // find import / export pairs
     let defs = DefAnalyzer::new(&output);
@@ -803,9 +953,10 @@ pub fn link(inputs: &mut [&mut rspirv::dr::Module], opts: &Options) -> Result<rs
     let matching_pairs = info.ensure_matching_import_export_pairs(&defs)?;
 
     // remove duplicates (https://github.com/KhronosGroup/SPIRV-Tools/blob/e7866de4b1dc2a7e8672867caeb0bdca49f458d3/source/opt/remove_duplicates_pass.cpp)
-    remove_duplicates(&mut output);
-
-    println!("{}\n\n", output.disassemble());
+    remove_duplicate_capablities(&mut output);
+    remove_duplicate_ext_inst_imports(&mut output);
+    let mut output = remove_duplicate_types(output);
+    // jb-todo: strip identical OpDecoration / OpDecorationGroups
 
     // remove names and decorations of import variables / functions https://github.com/KhronosGroup/SPIRV-Tools/blob/8a0ebd40f86d1f18ad42ea96c6ac53915076c3c7/source/opt/ir_context.cpp#L404
     import_kill_annotations_and_debug(&mut output, &info);
@@ -833,8 +984,6 @@ pub fn link(inputs: &mut [&mut rspirv::dr::Module], opts: &Options) -> Result<rs
         )],
     ));
 
-    println!("{}\n\n", output.disassemble());
-
     // output the module
     Ok(output)
 }

rspirv-linker/src/main.rs

@@ -1,5 +1,5 @@
-use rspirv_linker::*;
 use rspirv::binary::Disassemble;
+use rspirv_linker::*;
 
 fn main() -> Result<()> {
     let body1 = include_bytes!("../test/1/body_1.spv");

rspirv-linker/src/test.rs

@@ -368,11 +368,11 @@ mod test {
         let result = assemble_and_link(&[&a, &b], &Options::default())?;
 
         let expect = r#"OpModuleProcessed "Linked by rspirv-linker"
-            %1 = OpTypeFloat 32
-            %2 = OpTypeVoid
-            %3 = OpTypeFunction %2
-            %4 = OpVariable %1 Uniform
-            %5 = OpFunction %2 Inline %3
+            %1 = OpTypeVoid
+            %2 = OpTypeFloat 32
+            %3 = OpTypeFunction %1
+            %4 = OpVariable %2 Uniform
+            %5 = OpFunction %1 Inline %3
             %6 = OpLabel
             OpReturn
             OpFunctionEnd"#;
@@ -441,7 +441,7 @@ mod test {
         without_header_eq(result, expect);
         Ok(())
     }
-    /*
+
     #[test]
     fn names_and_decorations() -> Result<()> {
         let a = assemble_spirv(
@@ -489,29 +489,28 @@ mod test {
         let result = assemble_and_link(&[&a, &b], &Options::default())?;
 
         let expect = r#"OpCapability Kernel
-        OpName %1 "foo"
-        OpName %2 "param"
+        OpName %1 "param"
+        OpName %2 "foo"
         OpModuleProcessed "Linked by rspirv-linker"
+        OpDecorate %1 Restrict
         OpDecorate %3 Restrict
         OpDecorate %4 NonWritable
         %3 = OpDecorationGroup
         OpGroupDecorate %3 %4
-        OpDecorate %2 Restrict
         %5 = OpTypeVoid
         %6 = OpTypeInt 32 0
         %7 = OpTypePointer Function %6
         %8 = OpTypeFunction %5 %7
         %9 = OpFunction %5 None %8
         %4 = OpFunctionParameter %7
         OpFunctionEnd
-        %1 = OpFunction %5 None %8
-        %2 = OpFunctionParameter %7
+        %2 = OpFunction %5 None %8
+        %1 = OpFunctionParameter %7
         %10 = OpLabel
         OpReturn
-        OpFunctionEnd
-        "#;
+        OpFunctionEnd"#;
 
         without_header_eq(result, expect);
         Ok(())
-    }*/
+    }
 }

rustc_codegen_spirv/Cargo.toml

@@ -13,7 +13,7 @@ repository = "https://github.com/EmbarkStudios/rust-gpu"
 crate-type = ["dylib"]
 
 [dependencies]
-rspirv = { git = "https://github.com/gfx-rs/rspirv" }
+rspirv = "0.7.0"
 rspirv-linker = { path = "../rspirv-linker" }
 tar = "0.4"
 topological-sort = "0.1"