using IR level liveness info in register allocation

.gitignore

 target/*
 emit/*
 Cargo.lock
+*.log

src/compiler/backend/arch/x86_64/asm_backend.rs

@@ -15,6 +15,7 @@ use std::collections::HashMap;
 use std::str;
 use std::usize;
 use std::slice::Iter;
+use std::ops;
 
 struct ASMCode {
     name: MuTag, 
@@ -30,6 +31,10 @@ struct ASMCode {
     cond_branches: HashMap<usize, MuTag>,
     branches: HashMap<usize, MuTag>,
     
+    blocks: Vec<MuTag>,
+    block_start: HashMap<MuTag, usize>,
+    block_range: HashMap<MuTag, ops::Range<usize>>,
+    
     block_livein: HashMap<MuTag, Vec<MuID>>,
     block_liveout: HashMap<MuTag, Vec<MuID>>
 }
@@ -117,6 +122,25 @@ impl MachineCode for ASMCode {
         
         println!("");        
     }
+    
+    fn get_ir_block_livein(&self, block: MuTag) -> Option<&Vec<MuID>> {
+        self.block_livein.get(&block)
+    }
+    
+    fn get_ir_block_liveout(&self, block: MuTag) -> Option<&Vec<MuID>> {
+        self.block_liveout.get(&block)
+    }
+    
+    fn get_all_blocks(&self) -> &Vec<MuTag> {
+        &self.blocks
+    }
+    
+    fn get_block_range(&self, block: MuTag) -> Option<ops::Range<usize>> {
+        match self.block_range.get(&block) {
+            Some(r) => Some(r.clone()),
+            None => None
+        }
+    }
 }
 
 struct ASM {
@@ -426,6 +450,10 @@ impl CodeGenerator for ASMCodeGen {
                 cond_branches: HashMap::new(),
                 branches: HashMap::new(),
                 
+                blocks: vec![],
+                block_start: HashMap::new(),
+                block_range: HashMap::new(),
+                
                 block_livein: HashMap::new(),
                 block_liveout: HashMap::new()
             }));
@@ -464,6 +492,17 @@ impl CodeGenerator for ASMCodeGen {
     fn start_block(&mut self, block_name: MuTag) {
         let label = format!("{}:", self.asm_block_label(block_name));        
         self.add_asm_block_label(label, block_name);
+        self.cur_mut().blocks.push(block_name);
+        
+        let start = self.line();
+        self.cur_mut().block_start.insert(block_name, start);
+    }
+    
+    fn end_block(&mut self, block_name: MuTag) {
+        let start : usize = *self.cur().block_start.get(&block_name).unwrap();
+        let end : usize = self.line();
+        
+        self.cur_mut().block_range.insert(block_name, (start..end));
     }
     
     fn set_block_livein(&mut self, block_name: MuTag, live_in: &Vec<P<Value>>) {

src/compiler/backend/arch/x86_64/codegen.rs

@@ -13,6 +13,7 @@ pub trait CodeGenerator {
     fn start_block(&mut self, block_name: MuTag);
     fn set_block_livein(&mut self, block_name: MuTag, live_in: &Vec<P<Value>>);
     fn set_block_liveout(&mut self, block_name: MuTag, live_out: &Vec<P<Value>>);
+    fn end_block(&mut self, block_name: MuTag);
     
     fn emit_cmp_r64_r64(&mut self, op1: &P<Value>, op2: &P<Value>);
     fn emit_cmp_r64_imm32(&mut self, op1: &P<Value>, op2: u32);

src/compiler/backend/arch/x86_64/inst_sel.rs

@@ -380,7 +380,13 @@ impl <'a> InstructionSelection {
     }
     
     fn emit_common_prologue(&mut self, args: &Vec<P<Value>>) {
-        self.backend.start_block("prologue");
+        let block_name = "prologue";
+        self.backend.start_block(block_name);
+        
+        // no livein
+        // liveout = entry block's args
+        self.backend.set_block_livein(block_name, &vec![]);
+        self.backend.set_block_liveout(block_name, args);
         
         // push rbp
         self.backend.emit_push_r64(&x86_64::RBP);
@@ -414,10 +420,13 @@ impl <'a> InstructionSelection {
                 panic!("expect an arg value to be either int reg or fp reg");
             }
         }
+        
+        self.backend.end_block(block_name);
     }
     
     fn emit_common_epilogue(&mut self, ret_inst: &Instruction, cur_func: &MuFunction) {
-        self.backend.start_block("epilogue");
+        // epilogue is not a block (its a few instruction inserted before return)
+        // FIXME: this may change in the future
         
         // prepare return regs
         let ref ops = ret_inst.ops.borrow();
@@ -692,6 +701,8 @@ impl CompilerPass for InstructionSelection {
             for inst in block_content.body.iter() {
                 self.instruction_select(inst, func);
             }
+            
+            self.backend.end_block(block.label);
         }
     }
     

src/compiler/backend/reg_alloc/graph_coloring/liveness.rs

@@ -6,6 +6,7 @@ use ast::ir::*;
 use ast::types;
 use compiler::backend;
 use utils::vec_utils;
+use utils::LinkedHashSet;
 
 use std::collections::LinkedList;
 use std::collections::{HashMap, HashSet};
@@ -260,7 +261,111 @@ pub fn is_machine_reg(reg: MuID) -> bool {
     }
 }
 
+// from Tailoring Graph-coloring Register Allocation For Runtime Compilation, Figure 4
+pub fn build_chaitin_briggs (cf: &CompiledFunction, func: &MuFunction) -> InterferenceGraph {
+    let mut ig = InterferenceGraph::new();
+    
+    // precolor machine register nodes
+    for reg in backend::all_regs().iter() {
+        let reg_id = reg.extract_ssa_id().unwrap();
+        let node = ig.new_node(reg_id, &func.context);
+        ig.color_node(node, reg_id);
+    }
+    
+    // Initialize and creates nodes for all the involved temps/regs
+    for i in 0..cf.mc.number_of_insts() {
+        for reg_id in cf.mc.get_inst_reg_defines(i) {
+            let reg_id = *reg_id;
+            ig.new_node(reg_id, &func.context);
+        }
+        
+        for reg_id in cf.mc.get_inst_reg_uses(i) {
+            let reg_id = *reg_id;
+            ig.new_node(reg_id, &func.context);
+        }
+    }
+    
+    // all nodes has been added, we init graph (create adjacency matrix)
+    ig.init_graph();    
+    
+    for block in cf.mc.get_all_blocks() {
+        // Current_Live(B) = LiveOut(B)
+        let mut current_live = LinkedHashSet::from_vec(match cf.mc.get_ir_block_liveout(block) {
+            Some(liveout) => liveout.to_vec(),
+            None => panic!("cannot find liveout for block {}", block)
+        });
+        
+        let range = cf.mc.get_block_range(block);
+        if range.is_none() {
+            continue;
+        }
+        
+        // for every inst I in reverse order
+        for i in range.unwrap().rev() {
+            let src : Option<MuID> = {
+                if cf.mc.is_move(i) {
+                    let src = cf.mc.get_inst_reg_uses(i);
+                    let dst = cf.mc.get_inst_reg_defines(i);
+                    
+                    // src may be an immediate number
+                    // but dest is definitly a register
+                    debug_assert!(dst.len() == 1);
+                    
+                    if src.len() == 1 {
+                        let node1 = ig.get_node(src[0]);
+                        let node2 = ig.get_node(dst[0]);
+                        ig.add_move(node1, node2);
+                        
+                        Some(src[0])
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            };
+            
+            // for every definition D in I
+            for d in cf.mc.get_inst_reg_defines(i) {
+                // add an interference from D to every element E in Current_Live - {D}
+                // creating nodes if necessary
+                for e in current_live.iter() {
+                    if src.is_none() || (src.is_some() && *e != src.unwrap()) {
+                        let from = ig.get_node(*d);
+                        let to = ig.get_node(*e);
+                        
+                        if !ig.is_same_node(from, to) && !ig.is_adj(from, to) {
+                            if !ig.is_colored(from) {
+                                ig.add_interference_edge(from, to);
+                            }
+                            if !ig.is_colored(to) {
+                                ig.add_interference_edge(to, from);
+                            }
+                        }
+                    }
+                }
+            }
+            
+            // for every definition D in I
+            for d in cf.mc.get_inst_reg_defines(i) {
+                // remove D from Current_Live
+                current_live.remove(d);
+            }
+            
+            // for every use U in I
+            for u in cf.mc.get_inst_reg_uses(i) {
+                // add U to Current_live
+                current_live.insert(*u);
+            }
+        }
+    }
+    
+    ig
+}
+
 // from tony's code src/RegAlloc/Liveness.java
+// this function is no longer used
+#[allow(dead_code)]
 pub fn build (cf: &CompiledFunction, func: &MuFunction) -> InterferenceGraph {
     let mut ig = InterferenceGraph::new();
     

src/compiler/backend/reg_alloc/graph_coloring/mod.rs

@@ -2,5 +2,6 @@ mod liveness;
 mod coloring;
 
 pub use compiler::backend::reg_alloc::graph_coloring::liveness::InterferenceGraph;
-pub use compiler::backend::reg_alloc::graph_coloring::liveness::build as build_inteference_graph;
+//pub use compiler::backend::reg_alloc::graph_coloring::liveness::build as build_inteference_graph;
+pub use compiler::backend::reg_alloc::graph_coloring::liveness::build_chaitin_briggs as build_inteference_graph;
 pub use compiler::backend::reg_alloc::graph_coloring::coloring::GraphColoring;
\ No newline at end of file

src/lib.rs

@@ -4,7 +4,7 @@ extern crate lazy_static;
 extern crate log;
 
 #[macro_use]
-mod utils;
+pub mod utils;
 pub mod ast;
 pub mod vm;
 pub mod compiler;

src/utils/linked_hashset.rs

@@ -13,6 +13,16 @@ impl<K: Hash + Eq> LinkedHashSet<K> {
     pub fn new() -> Self {
         LinkedHashSet(LinkedHashMap::new())
     }
+    
+    pub fn from_vec(from: Vec<K>) -> Self {
+        let mut ret = LinkedHashSet::new();
+        
+        for ele in from {
+            ret.insert(ele);
+        }
+        
+        ret
+    }
 }
 
 impl<K: Hash + Eq, S: BuildHasher> LinkedHashSet<K, S> {

src/utils/mod.rs

@@ -20,6 +20,18 @@ macro_rules! select_value {
 pub mod vec_utils {
     use std::fmt;
     
+    pub fn is_identical_to_str_ignore_order<T: Ord + fmt::Display + Clone, Q: Ord + fmt::Display + Clone> (vec: &Vec<T>, mut expect: Vec<Q>) -> bool {
+        let mut vec_copy = vec.to_vec();
+        vec_copy.sort();
+        
+        expect.sort();
+        
+        let a = as_str(&vec_copy);
+        let b = as_str(&expect);
+        
+        a == b
+    }
+    
     pub fn as_str<T: fmt::Display>(vec: &Vec<T>) -> String {
         let mut ret = String::new();
         for i in 0..vec.len() {

src/vm/machine_code.rs

 use ast::ir::*;
+use std::ops;
 
 pub struct CompiledFunction {
     pub fn_name: MuTag,
@@ -19,4 +20,10 @@ pub trait MachineCode {
     fn get_inst_reg_defines(&self, index: usize) -> &Vec<MuID>;
     
     fn replace_reg(&mut self, from: MuID, to: MuID);
+    
+    fn get_ir_block_livein(&self, block: MuTag) -> Option<&Vec<MuID>>;
+    fn get_ir_block_liveout(&self, block: MuTag) -> Option<&Vec<MuID>>;
+    
+    fn get_all_blocks(&self) -> &Vec<MuTag>;
+    fn get_block_range(&self, block: MuTag) -> Option<ops::Range<usize>>;
 }
\ No newline at end of file

tests/test_compiler/test_regalloc.rs

@@ -4,9 +4,63 @@ extern crate simple_logger;
 
 use test_ir::test_ir::factorial;
 use self::mu::compiler::*;
+use self::mu::utils::vec_utils;
+use self::mu::ast::ir::*;
 
 use std::sync::Arc;
 
+#[test]
+fn test_ir_liveness_fac() {
+    simple_logger::init_with_level(log::LogLevel::Trace).ok();
+    
+    let vm_context = Arc::new(factorial());
+    
+    let compiler = Compiler::new(CompilerPolicy::new(vec![
+            Box::new(passes::DefUse::new()),
+            Box::new(passes::TreeGen::new()),
+            Box::new(passes::ControlFlowAnalysis::new()),
+            Box::new(passes::TraceGen::new()),
+            Box::new(backend::inst_sel::InstructionSelection::new()),
+    ]), vm_context.clone());
+    
+    let funcs = vm_context.funcs().read().unwrap();
+    let mut factorial_func = funcs.get("fac").unwrap().borrow_mut();
+    
+    compiler.compile(&mut factorial_func);
+    
+    let cf_lock = vm_context.compiled_funcs().read().unwrap();
+    let cf = cf_lock.get("fac").unwrap().borrow();
+    
+    // block 0
+    
+    let block_0_livein = cf.mc.get_ir_block_livein("blk_0").unwrap();
+    let blk_0_n_3 = factorial_func.context.get_value_by_tag("blk_0_n_3").unwrap().id;
+    assert!(vec_utils::is_identical_to_str_ignore_order(block_0_livein, vec![blk_0_n_3]));
+    
+    let block_0_liveout = cf.mc.get_ir_block_liveout("blk_0").unwrap();
+    assert!(vec_utils::is_identical_to_str_ignore_order(block_0_liveout, vec![blk_0_n_3]));
+    
+    // block 1
+    
+    let block_1_livein = cf.mc.get_ir_block_livein("blk_1").unwrap();
+    let blk_1_n_3 = factorial_func.context.get_value_by_tag("blk_1_n_3").unwrap().id;
+    assert!(vec_utils::is_identical_to_str_ignore_order(block_1_livein, vec![blk_1_n_3]));
+    
+    let block_1_liveout = cf.mc.get_ir_block_liveout("blk_1").unwrap();
+    let blk_1_v52 = factorial_func.context.get_value_by_tag("blk_1_v52").unwrap().id;
+    assert!(vec_utils::is_identical_to_str_ignore_order(block_1_liveout, vec![blk_1_v52]));
+    
+    // block 2
+    
+    let block_2_livein = cf.mc.get_ir_block_livein("blk_2").unwrap();
+    let blk_2_v53 = factorial_func.context.get_value_by_tag("blk_2_v53").unwrap().id;
+    assert!(vec_utils::is_identical_to_str_ignore_order(block_2_livein, vec![blk_2_v53]));
+    
+    let block_2_liveout = cf.mc.get_ir_block_liveout("blk_2").unwrap();
+    let expect : Vec<MuID> = vec![];
+    assert!(vec_utils::is_identical_to_str_ignore_order(block_2_liveout, expect));
+}
+
 #[test]
 fn test_regalloc_fac() {
     simple_logger::init_with_level(log::LogLevel::Trace).ok();

tests/test_ir/test_ir.rs

@@ -275,6 +275,7 @@ pub fn factorial() -> VMContext {
         v: Instruction_::BinOp(BinOp::Mul, 0, 1)
     });
 
+    // BRANCH blk_2 (%blk_1_v52)
     let blk_1_term = func.new_inst(Instruction{
         value: None,
         ops: RefCell::new(vec![blk_1_v52.clone()]),