fix: support load store indirection

Author: molikto

crates/rustc_codegen_spirv/src/linker/inline_globals.rs

@@ -1,13 +1,80 @@
 use rspirv::binary::Disassemble;
 use rspirv::dr::{Instruction, Module, Operand};
-use rspirv::spirv::Op;
+use rspirv::spirv::{Op, StorageClass};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_session::Session;
 
+// bool is if this needs stored
+#[derive(Debug, Clone, PartialEq)]
+struct NormalizedInstructions {
+    vars: Vec<Instruction>,
+    insts: Vec<Instruction>,
+    root: u32,
+}
+
+impl NormalizedInstructions {
+    fn new(id: u32) -> Self {
+        NormalizedInstructions {
+            vars: Vec::new(),
+            insts: Vec::new(),
+            root: id,
+        }
+    }
+
+    fn extend(&mut self, o: NormalizedInstructions) {
+        self.vars.extend(o.vars);
+        self.insts.extend(o.insts);
+    }
+
+    fn is_empty(&self) -> bool {
+        self.insts.is_empty() && self.vars.is_empty()
+    }
+
+    fn fix_ids(&mut self, bound: &mut u32, new_root: u32) {
+        let mut id_map: FxHashMap<u32, u32> = FxHashMap::default();
+        id_map.insert(self.root, new_root);
+        for inst in &mut self.vars {
+            Self::fix_instruction(self.root, inst, &mut id_map, bound, new_root);
+        }
+        for inst in &mut self.insts {
+            Self::fix_instruction(self.root, inst, &mut id_map, bound, new_root);
+        }
+    }
+
+    fn fix_instruction(
+        root: u32,
+        inst: &mut Instruction,
+        id_map: &mut FxHashMap<u32, u32>,
+        bound: &mut u32,
+        new_root: u32,
+    ) {
+            for op in &mut inst.operands {
+                match op {
+                    Operand::IdRef(id) => match id_map.get(id) {
+                        Some(new_id) => {
+                            *id = *new_id;
+                        }
+                        _ => {}
+                    },
+                    _ => {}
+                }
+            }
+        if let Some(id) = &mut inst.result_id {
+            if *id != root {
+                id_map.insert(*id, *bound);
+                *id = *bound;
+                *bound += 1;
+            } else {
+                *id = new_root;
+            }
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq)]
 enum FunctionArg {
     Invalid,
-    Insts(Vec<Instruction>),
+    Insts(NormalizedInstructions),
 }
 
 pub fn inline_global_varaibles(sess: &Session, module: &mut Module) -> super::Result<()> {
@@ -36,14 +103,30 @@ fn inline_global_varaibles_rec(sess: &Session, module: &mut Module) -> super::Re
     }
     // then we keep track of which function parameter are always called with the same expression that only uses global variables
     let mut function_args: FxHashMap<(u32, u32), FunctionArg> = FxHashMap::default();
+    let mut bound = module.header.as_ref().unwrap().bound;
     for caller in &module.functions {
         let mut insts: FxHashMap<u32, Instruction> = FxHashMap::default();
+        // for variables that only stored once and it's stored as a ref
+        let mut ref_stores: FxHashMap<u32, Option<u32>> = FxHashMap::default();
         for block in &caller.blocks {
             for inst in &block.instructions {
                 if inst.result_id.is_some() {
                     insts.insert(inst.result_id.unwrap(), inst.clone());
                 }
-                if inst.class.opcode == Op::FunctionCall {
+                if inst.class.opcode == Op::Store {
+                    if let Operand::IdRef(to) = inst.operands[0] {
+                        if let Operand::IdRef(from) = inst.operands[1] {
+                            match ref_stores.get(&to) {
+                                None => {
+                                    ref_stores.insert(to, Some(from));
+                                }
+                                Some(_) => {
+                                    ref_stores.insert(to, None);
+                                }
+                            }
+                        }
+                    }
+                } else if inst.class.opcode == Op::FunctionCall {
                     let function_id = match &inst.operands[0] {
                         &Operand::IdRef(w) => w,
                         _ => panic!(),
@@ -52,16 +135,19 @@ fn inline_global_varaibles_rec(sess: &Session, module: &mut Module) -> super::Re
                         let key = (function_id, i as u32 - 1);
                         match &inst.operands[i] {
                             &Operand::IdRef(w) => match &function_args.get(&key) {
-                                None => match get_const_arg_insts(&variables, &insts, w) {
-                                    Some(insts) => {
-                                        function_args.insert(key, FunctionArg::Insts(insts));
-                                    }
-                                    None => {
-                                        function_args.insert(key, FunctionArg::Invalid);
+                                None => {
+                                    match get_const_arg_insts(bound, &variables, &insts, &ref_stores, w) {
+                                        Some(insts) => {
+                                            function_args.insert(key, FunctionArg::Insts(insts));
+                                        }
+                                        None => {
+                                            function_args.insert(key, FunctionArg::Invalid);
+                                        }
                                     }
-                                },
+                                }
                                 Some(FunctionArg::Insts(w2)) => {
-                                    let new_insts = get_const_arg_insts(&variables, &insts, w);
+                                    let new_insts =
+                                        get_const_arg_insts(bound, &variables, &insts, &ref_stores, w);
                                     match new_insts {
                                         Some(new_insts) => {
                                             if new_insts != *w2 {
@@ -93,11 +179,10 @@ fn inline_global_varaibles_rec(sess: &Session, module: &mut Module) -> super::Re
     if function_args.is_empty() {
         return Ok(false);
     }
-    let mut bound = module.header.as_ref().unwrap().bound;
     for function in &mut module.functions {
         let def = function.def.as_mut().unwrap();
         let fid = def.result_id.unwrap();
-        let mut insts: Vec<Instruction> = Vec::new();
+        let mut insts = NormalizedInstructions::new(0);
         let mut j: u32 = 0;
         let mut i = 0;
         let mut removed_indexes: Vec<u32> = Vec::new();
@@ -108,10 +193,7 @@ fn inline_global_varaibles_rec(sess: &Session, module: &mut Module) -> super::Re
                 Some(FunctionArg::Insts(arg)) => {
                     let parameter = function.parameters.remove(i);
                     let mut arg = arg.clone();
-                    arg.reverse();
-                    insts_replacing_captured_ids(&mut arg, &mut bound);
-                    let index = arg.len() - 1;
-                    arg[index].result_id = parameter.result_id;
+                    arg.fix_ids(&mut bound, parameter.result_id.unwrap());
                     insts.extend(arg);
                     removed_indexes.push(j);
                     removed = true;
@@ -132,15 +214,16 @@ fn inline_global_varaibles_rec(sess: &Session, module: &mut Module) -> super::Re
                 for i in removed_indexes.iter().rev() {
                     let i = *i as usize + 1;
                     function_type.operands.remove(i);
-                    function_type.result_id = Some(tid);
                 }
+                function_type.result_id = Some(tid);
                 def.operands[1] = Operand::IdRef(tid);
                 module.types_global_values.push(function_type);
             }
         }
         // callee side. insert initialization instructions, which reuse the ids of the removed parameters
         if !function.blocks.is_empty() {
             let first_block = &mut function.blocks[0];
+            first_block.instructions.splice(0..0, insts.vars);
             // skip some instructions that must be at top of block
             let mut i = 0;
             loop {
@@ -154,7 +237,7 @@ fn inline_global_varaibles_rec(sess: &Session, module: &mut Module) -> super::Re
                 }
                 i += 1;
             }
-            first_block.instructions.splice(i..i, insts);
+            first_block.instructions.splice(i..i, insts.insts);
         }
         // caller side, remove parameters from function call
         for block in &mut function.blocks {
@@ -181,76 +264,90 @@ fn inline_global_varaibles_rec(sess: &Session, module: &mut Module) -> super::Re
     Ok(true)
 }
 
-fn insts_replacing_captured_ids(arg: &mut Vec<Instruction>, bound: &mut u32) {
-    let mut id_map: FxHashMap<u32, u32> = FxHashMap::default();
-    for ins in arg {
-        if let Some(id) = &mut ins.result_id {
-            for op in &mut ins.operands {
-                match op {
-                    Operand::IdRef(id) => match id_map.get(id) {
-                        Some(new_id) => {
-                            *id = *new_id;
-                        }
-                        _ => {}
-                    },
-                    _ => {}
-                }
-            }
-            id_map.insert(*id, *bound);
-            *id = *bound;
-            *bound += 1;
-        }
-    }
-}
-
 fn get_const_arg_operands(
     variables: &FxHashSet<u32>,
     insts: &FxHashMap<u32, Instruction>,
+    ref_stores: &FxHashMap<u32, Option<u32>>,
     operand: &Operand,
-) -> Option<Vec<Instruction>> {
+) -> Option<NormalizedInstructions> {
     match operand {
         Operand::IdRef(id) => {
-            let insts = get_const_arg_insts(variables, insts, *id)?;
+            let insts = get_const_arg_insts_rec(variables, insts, ref_stores, *id)?;
             return Some(insts);
         }
-        Operand::LiteralInt32(_) => {},
-        Operand::LiteralInt64(_) => {},
-        Operand::LiteralFloat32(_) => {},
-        Operand::LiteralFloat64(_) => {},
-        Operand::LiteralExtInstInteger(_) => {},
-        Operand::LiteralSpecConstantOpInteger(_) => {},
-        Operand::LiteralString(_) => {},
+        Operand::LiteralInt32(_) => {}
+        Operand::LiteralInt64(_) => {}
+        Operand::LiteralFloat32(_) => {}
+        Operand::LiteralFloat64(_) => {}
+        Operand::LiteralExtInstInteger(_) => {}
+        Operand::LiteralSpecConstantOpInteger(_) => {}
+        Operand::LiteralString(_) => {}
         _ => {
             // TOOD add more cases
             return None;
         }
     }
-    return Some(Vec::new());
+    return Some(NormalizedInstructions::new(0));
 }
 
 fn get_const_arg_insts(
+    mut bound: u32,
+    variables: &FxHashSet<u32>,
+    insts: &FxHashMap<u32, Instruction>,
+    ref_stores: &FxHashMap<u32, Option<u32>>,
+    id: u32,
+) -> Option<NormalizedInstructions> {
+    let mut res = get_const_arg_insts_rec(variables, insts, ref_stores, id)?;
+    res.insts.reverse();
+    // the bound passed in is always the same
+    // we need to normalize the ids, so they are the same when compared
+    let fake_root = bound;
+    bound += 1;
+    res.fix_ids(&mut bound, fake_root);
+    res.root = fake_root;
+    Some(res)
+}
+
+fn get_const_arg_insts_rec(
     variables: &FxHashSet<u32>,
     insts: &FxHashMap<u32, Instruction>,
+    ref_stores: &FxHashMap<u32, Option<u32>>,
     id: u32,
-) -> Option<Vec<Instruction>> {
-    let mut result: Vec<Instruction> = Vec::new();
+) -> Option<NormalizedInstructions> {
+    let mut result = NormalizedInstructions::new(id);
     if variables.contains(&id) {
         return Some(result);
     }
     let par: &Instruction = insts.get(&id)?;
     if par.class.opcode == Op::AccessChain {
-        result.push(par.clone());
+        result.insts.push(par.clone());
         for oprand in &par.operands {
-            let insts = get_const_arg_operands(variables, insts, oprand)?;
+            let insts = get_const_arg_operands(variables, insts, ref_stores, oprand)?;
             result.extend(insts);
         }
     } else if par.class.opcode == Op::FunctionCall {
-        result.push(par.clone());
+        result.insts.push(par.clone());
         // skip first, first is function id
         for oprand in &par.operands[1..] {
-            let insts = get_const_arg_operands(variables, insts, oprand)?;
+            let insts = get_const_arg_operands(variables, insts, ref_stores, oprand)?;
             result.extend(insts);
         }
+    } else if par.class.opcode == Op::Variable {
+        result.vars.push(par.clone());
+        let stored = ref_stores.get(&id)?;
+        let stored = (*stored)?;
+        result.insts.push(Instruction::new(
+            Op::Store,
+            None,
+            None,
+            vec![Operand::IdRef(id), Operand::IdRef(stored)],
+        ));
+        let new_insts = get_const_arg_insts_rec(variables, insts, ref_stores, stored)?;
+        result.extend(new_insts);
+    } else if par.class.opcode == Op::ArrayLength {
+        result.insts.push(par.clone());
+        let insts = get_const_arg_operands(variables, insts, ref_stores, &par.operands[0])?;
+        result.extend(insts);
     } else {
         // TOOD add more cases
         return None;

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -220,19 +220,20 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
     }
 
     // this is needed so we can inline more global variables
-    {
-        let _timer = sess.timer("simpl_op_store_var");
-        simpl_op_store_var::simpl_op_store_var(sess, &mut output)?;
-    }
+    // {
+    //     let _timer = sess.timer("simpl_op_store_var");
+    //     simpl_op_store_var::simpl_op_store_var(sess, &mut output)?;
+    // }
 
     {
         let _timer = sess.timer("link_inline_global");
         inline_globals::inline_global_varaibles(sess, &mut output)?;
     }
 
+    // needed because inline global create duplicate types...
     {
-        let _timer = sess.timer("link_inline_global");
-        inline_globals::inline_global_varaibles(sess, &mut output)?;
+        let _timer = sess.timer("link_remove_duplicate_types_round_2");
+        duplicates::remove_duplicate_types(&mut output);
     }
 
     {
@@ -320,8 +321,8 @@ pub fn link(sess: &Session, mut inputs: Vec<Module>, opts: &Options) -> Result<L
         let _timer = sess.timer("link_sort_globals");
         simple_passes::sort_globals(&mut output);
     }
-
-    std::fs::write("res.txt", output.disassemble());
+    
+    // std::fs::write("res.txt", output.disassemble());
 
     let mut output = if opts.emit_multiple_modules {
         let modules = output