Merge pull request #33 from ethanuppal/regalloc

Register Allocation

.ocamlinit

@@ -3,3 +3,14 @@ open X86ISTMB
 #install_printer Ast.pp_expr
 #install_printer Ast.pp_stmt
 #install_printer Ast.pp_prog
+#install_printer Variable.pp
+#install_printer Id.pp
+
+let show_regalloc file =
+  let source = Util.read_file file in
+  let prog = Parse_lex.lex_and_parse source in
+  let cfg = Ir_gen.generate prog |> List.hd in
+  let liveliness = Liveliness.analysis_of cfg in
+  let ordering = InstrOrdering.make cfg in
+  print_endline (Cfg.to_string cfg);
+  Regalloc.allocate_for cfg liveliness ordering |> Regalloc.VarTbl.to_seq |> List.of_seq

README.md

@@ -7,7 +7,7 @@
 ![CI Status](https://github.com/ethanuppal/cs3110_compiler/actions/workflows/ci.yaml/badge.svg)
 
 > "x86 is simple trust me bro"  
-> Last updated: 2024-05-09 01:08:34.998129
+> Last updated: 2024-05-10 22:08:20.895180
 
 ```
 $ ./main -h

lib/backend/asm.ml

@@ -38,6 +38,8 @@ module Register = struct
     | R13 -> "r13"
     | R14 -> "r14"
     | R15 -> "r15"
+
+  let compare = Stdlib.compare
 end
 
 module Operand = struct

lib/backend/liveliness.ml

@@ -136,7 +136,9 @@ let apply_rules liveliness analysis cfg bb ir ir_idx ~is_final =
   | Add (var, op1, op2) | Sub (var, op1, op2) | TestEqual (var, op1, op2) ->
       write_var var;
       read_op op1;
-      read_op op2);
+      read_op op2
+  | Call (var, _) -> write_var var
+  | Return op -> read_op op);
   check_for_changes ()
 
 (** [pass work_list liveliness cfg bb] performs a single pass of liveliness

lib/backend/regalloc/instrOrdering.ml

@@ -0,0 +1,26 @@
+module OrderMap = Hashtbl.Make (Id)
+
+type t = int OrderMap.t
+type instr_id = Id.t * int
+
+let make cfg =
+  (* arbitrary order for now, can be improved later. does not actually matter,
+     any ordering works. better orderings only improve performance. *)
+  let map = OrderMap.create 10 in
+  let n = ref 0 in
+  Cfg.iter
+    (fun bb ->
+      OrderMap.add map (Basic_block.id_of bb) !n;
+      n := !n + 1)
+    cfg;
+  map
+
+let compare (ordering : t) (bb_id_1, instr_idx_1) (bb_id_2, instr_idx_2) =
+  let bb_order_1 = OrderMap.find ordering bb_id_1 in
+  let bb_order_2 = OrderMap.find ordering bb_id_2 in
+  if bb_order_1 > bb_order_2 then 1
+  else if bb_order_1 = bb_order_2 then
+    if instr_idx_1 > instr_idx_2 then 1
+    else if instr_idx_1 = instr_idx_2 then 0
+    else -1
+  else -1

lib/backend/regalloc/instrOrdering.mli

@@ -0,0 +1,14 @@
+type t
+type instr_id = Id.t * int
+
+(** [make cfg] is an ordering of the instructions in [cfg] in chronological
+    order. *)
+val make : Cfg.t -> t
+
+(** [compare ordering id1 id2] is 0 if [id1 = id2], negative if the instruction
+    with id [id1] comes before the instruction with id [id2], and positive if it
+    comes after.
+
+    Requires that [id1] and [id2] are valid ids for the cfg this ordering was
+    constructed from. *)
+val compare : t -> instr_id -> instr_id -> int

lib/backend/regalloc/regalloc.ml

@@ -0,0 +1,134 @@
+open Util
+module VarTbl = Hashtbl.Make (Variable)
+
+(* TODO: standardize instruction id? *)
+type instr_id = Id.t * int
+
+(** [start] is the first instruction (in terms of the arbitrary ordering)
+    *after* which a variable is live. [stop] is the last instruction *after*
+    which a variable is live.
+
+    See definition on page 898: https://dl.acm.org/doi/pdf/10.1145/330249.330250 *)
+type interval = {
+  start : instr_id;
+  stop : instr_id;
+}
+
+type allocation =
+  | Register of Asm.Register.t
+  | Spill
+
+module BBAnalysis = Liveliness.BasicBlockAnalysis
+
+let registers =
+  let open Asm.Register in
+  [ RAX; RBX; RCX; RDX; RBP; RSI; RDI; R8; R9; R10; R11; R12; R13; R14; R15 ]
+
+let live_intervals (cfg : Cfg.t) (liveliness : BBAnalysis.t IdMap.t)
+    (ordering : InstrOrdering.t) =
+  let tbl = VarTbl.create 16 in
+
+  let expand_interval original live_id =
+    let cmp = InstrOrdering.compare ordering in
+    if cmp live_id original.start < 0 then { original with start = live_id }
+    else if cmp live_id original.stop > 0 then { original with stop = live_id }
+    else original
+  in
+
+  let update_table instr_id used_set =
+    Liveliness.VariableSet.iter
+      (fun live ->
+        let current_opt = VarTbl.find_opt tbl live in
+        let new_interval =
+          match current_opt with
+          | None -> { start = instr_id; stop = instr_id }
+          | Some current -> expand_interval current instr_id
+        in
+        VarTbl.replace tbl live new_interval)
+      used_set
+  in
+
+  Cfg.iter
+    (fun bb ->
+      let bb_id = Basic_block.id_of bb in
+      let analysis = IdMap.find liveliness bb_id in
+      for instr_idx = 0 to Basic_block.length_of bb - 1 do
+        let live_set = BBAnalysis.live_after_instr analysis instr_idx in
+        let kill_var = Basic_block.get_ir bb instr_idx |> Ir.kill_of in
+        let used_set =
+          match kill_var with
+          | Some var -> Liveliness.VariableSet.add var live_set
+          | None -> live_set
+        in
+        update_table (bb_id, instr_idx) used_set
+      done)
+    cfg;
+
+  VarTbl.to_seq tbl |> List.of_seq
+
+(* Algorithm source:
+   https://en.wikipedia.org/wiki/Register_allocation#Pseudocode *)
+let linear_scan (intervals : (Variable.t * interval) list)
+    (ordering : InstrOrdering.t) =
+  let compare_instr_id = InstrOrdering.compare ordering in
+  let compare_pair_start (_, i1) (_, i2) = compare_instr_id i1.start i2.start in
+  let compare_pair_end (_, i1) (_, i2) = compare_instr_id i1.stop i2.stop in
+  let sorted_intervals = List.sort compare_pair_start intervals in
+
+  let assigned_alloc : allocation VarTbl.t = VarTbl.create 4 in
+
+  let module RegSet = Set.Make (Asm.Register) in
+  let free_registers : RegSet.t ref = ref (RegSet.of_list registers) in
+
+  (* must remain sorted by increasing end point *)
+  let active : (Variable.t * interval) BatRefList.t = BatRefList.empty () in
+
+  let expire_old_intervals (current : interval) =
+    (* this is also really annoying because BatRefList has no partition *)
+    BatRefList.filter
+      (fun (var, interval) ->
+        let keep = compare_instr_id interval.stop current.start >= 0 in
+        (if not keep then
+           let alloc = VarTbl.find assigned_alloc var in
+           match alloc with
+           | Register r -> free_registers := RegSet.add r !free_registers
+           | Spill -> failwith "Interval in active cannot be spilled");
+        keep)
+      active
+  in
+
+  let spill_at_interval ((var, interval) : Variable.t * interval) =
+    let spill_var, spill_interval = BatRefList.last active in
+
+    if compare_instr_id spill_interval.stop interval.stop > 0 then (
+      (* spill guaranteed to be assigned an actual register *)
+      let alloc = VarTbl.find assigned_alloc spill_var in
+      VarTbl.replace assigned_alloc var alloc;
+      VarTbl.replace assigned_alloc spill_var Spill;
+
+      (* this sucks. can we maybe keep active in reverse order? *)
+      BatRefList.Index.remove_at active (BatRefList.length active - 1);
+
+      (* add_sort is buggy... TODO: new impl *)
+      BatRefList.push active (var, interval);
+      BatRefList.sort ~cmp:compare_pair_end active)
+    else VarTbl.replace assigned_alloc var Spill
+  in
+
+  List.iter
+    (fun (var, interval) ->
+      expire_old_intervals interval;
+      match RegSet.choose_opt !free_registers with
+      | Some register ->
+          free_registers := RegSet.remove register !free_registers;
+          VarTbl.replace assigned_alloc var (Register register);
+          BatRefList.push active (var, interval);
+          BatRefList.sort ~cmp:compare_pair_end active
+      | None -> spill_at_interval (var, interval))
+    sorted_intervals;
+
+  assigned_alloc
+
+let allocate_for cfg liveliness ordering =
+  let vars_with_intervals = live_intervals cfg liveliness ordering in
+  linear_scan vars_with_intervals ordering

lib/backend/regalloc/regalloc.mli

@@ -0,0 +1,14 @@
+open Util
+module VarTbl : Hashtbl.S with type key = Variable.t
+
+type allocation =
+  | Register of Asm.Register.t
+  | Spill
+
+val registers : Asm.Register.t list
+
+val allocate_for :
+  Cfg.t ->
+  Liveliness.BasicBlockAnalysis.t IdMap.t ->
+  InstrOrdering.t ->
+  allocation VarTbl.t

lib/core/util.ml

@@ -76,6 +76,7 @@ module ArrayView : sig
   val length : 'a t -> int
   val get : 'a t -> int -> 'a
   val last : 'a t -> 'a
+  val iteri : (int -> 'a -> unit) -> 'a t -> unit
 end = struct
   let from_bat_dyn_arr = id
 

lib/ir/id.ml

@@ -5,6 +5,7 @@ let equal = Int.equal
 let hash = Int.hash
 let compare = Int.compare
 let int_of id = id
+let pp fmt id = Format.fprintf fmt "id[%i]" id
 
 module Gen = struct
   type t = {

lib/ir/id.mli

@@ -20,6 +20,8 @@ val compare : id -> id -> int
 (** [hash id] is a hash of [id]. *)
 val hash : id -> int
 
+val pp : Format.formatter -> id -> unit
+
 module Gen : sig
   (** Values of type [t] are unique-identifier generators. *)
   type t

lib/ir/ir.ml

@@ -9,7 +9,8 @@ type t =
   | Deref of Variable.t * Operand.t
   | TestEqual of Variable.t * Operand.t * Operand.t
   | DebugPrint of Operand.t
-(* | Call of string * Operand.t list *)
+  | Call of Variable.t * string
+  | Return of Operand.t
 
 (** [kill_of ir] is [Some var] if [var] is assigned to in [ir] and [None]
     otherwise. *)
@@ -20,7 +21,7 @@ let kill_of = function
   | Ref (var, _)
   | Deref (var, _)
   | TestEqual (var, _, _) -> Some var
-  | DebugPrint _ -> None
+  | DebugPrint _ | Call _ | Return _ -> None
 
 let to_string =
   let open Printf in
@@ -41,3 +42,5 @@ let to_string =
       sprintf "%s = %s == %s" (Variable.to_string r) (Operand.to_string o1)
         (Operand.to_string o2)
   | DebugPrint op -> sprintf "debug_print %s" (Operand.to_string op)
+  | Call (r, name) -> sprintf "%s = %s()" (Variable.to_string r) name
+  | Return op -> sprintf "return %s" (Operand.to_string op)

lib/ir/ir_sim.ml

@@ -33,7 +33,9 @@ let run simulator cfg =
                failwith "Ir_sim does not support pointers"
            | Ir.DebugPrint oper ->
                simulator.output <-
-                 simulator.output ^ Printf.sprintf "%d\n" (eval oper));
+                 simulator.output ^ Printf.sprintf "%d\n" (eval oper)
+           | Ir.Call _ | Ir.Return _ ->
+               failwith "Ir_sim does not support function calls yet");
     let cond =
       match Basic_block.condition_of bb with
       | Always -> true

lib/ir/variable.ml

@@ -6,6 +6,7 @@ let var_gen = Id.Gen.make ()
 let make () = Id.Gen.next var_gen
 let id_of var = var
 let to_string = Id.int_of >> string_of_int >> ( ^ ) "i"
+let pp = pp_of to_string
 let compare = Id.compare
 let equal = Id.equal
 let hash = Id.hash

lib/ir/variable.mli

@@ -10,6 +10,7 @@ val id_of : t -> Id.t
 (** [to_string var] is [var] as a string. *)
 val to_string : t -> string
 
+val pp : Format.formatter -> t -> unit
 val compare : t -> t -> int
 val equal : t -> t -> bool
 val hash : t -> int

source/regalloc_spill_test.x86istmb

@@ -0,0 +1,36 @@
+func main() {
+    let a = 0
+    let b = 1
+    let c = 2
+    let d = 3
+    let e = 4
+    let f = 5
+    let g = 6
+    let h = 7
+    let i = 8
+    let j = 9
+    let k = 10
+    let l = 11
+    let m = 12
+    let n = 13
+    let o = 14
+    let q = 16
+    let p = 15
+    print a
+    print b
+    print c
+    print d
+    print e
+    print f
+    print g
+    print h
+    print i
+    print j
+    print k
+    print l
+    print m
+    print n
+    print o
+    print p
+    print q
+}

source/regalloc_test.x86istmb

@@ -0,0 +1,16 @@
+func main() {
+    let x = 0
+    let y = 2
+    let z = x
+    let w = z + y
+    z = 10
+    if (true) {
+        let x = 8
+        let y = w
+        z = 1
+        let a = x + y
+    }
+    let a = 9
+    print w
+    print a
+}