start register allocation

1 files changed, 131 insertions, 7 deletions

diff --git a/lo.ml b/lo.ml
index 61df4b1..6127127 100644
--- a/lo.ml
+++ b/lo.ml
@@ -1,7 +1,7 @@
 type id = int
 module ISet = Set.Make
   (struct
-    type t = id
+    type t = int
     let compare = compare
   end)
 
@@ -97,8 +97,95 @@ let liveness p =
   done;
   liveout
 
+type reginfo =
+  { mutable rreg: int
+  ; mutable rspill: int option
+  ; mutable rhint: int
+  }
+
+let regalloc nr p l =
+  let regs = Array.init (Array.length p)
+    (fun _ ->
+      { rreg = -1
+      ; rspill = None
+      ; rhint = -1
+      }) in
+
+  let ( |> ) a b = if a = -1 then b else a in
+
+  (* Number of spilled registers. *)
+  let spill = ref 0 in
+  let rspill i =
+    if regs.(i).rspill = None then begin
+      regs.(i).rspill <- Some !spill;
+      incr spill;
+    end in
 
+  (* Associative list binding irrefs to registers,
+     it is ordered by freshness. *)
+  let used = ref [] in
+  let free = ref (
+      let rec m i = if i = nr then []  else i :: m (i+1)
+      in m 0
+    ) in
+  let alloc hint i =
+    let r, fl =
+      let l, fl = List.partition ((=) hint) !free in
+      if l <> [] then (hint, fl) else
+      match !free with
+      | r :: fl -> (r, fl)
+      | [] ->
+        (* No more free registers, we need to spill. *)
+        let rec g = function
+          | [] -> assert false
+          | [r,i'] -> rspill i'; (r, [])
+          | x :: us ->
+            let (r, us) = g us in
+            (r, x :: us) in
+        let r, us = g !used in
+        used := us;
+        r, [] in
+    free := fl;
+    used := (r, i) :: !used;
+    r in
+
+  for i = 0 to Array.length p -1 do
+
+    (* Forget about all bindings that are not live
+       at the end of the instruction. *)
+    let used', free' = List.partition
+      (fun (_, i') -> ISet.mem i' l.(i)) !used in
+    used := used';
+    free := List.map fst free' @ !free;
+
+    (* Bind a register to the current instruction
+       if its result is not discarded. *)
+    if ISet.mem i l.(i) then begin
+      match p.(i) with
+      | ICon _ | IBrz _ | IJmp _ | INop -> ()
+      | IPhi l ->
+        (* Try to ensure that variables merged by a phi
+           use the same register. *)
+        let f r {pvar;_} = regs.(pvar).rreg |> r in
+        let r = List.fold_left f (-1) l in
+        let r =
+          let h = regs.(i).rhint in
+          if r = -1 then alloc h i else r in
+        List.iter (fun {pvar;_} ->
+          regs.(pvar).rhint <- r
+        ) l;
+        regs.(i).rreg <- r
+      | IUop (_, i')
+      | IBop (i', _, _) ->
+        let h =
+          regs.(i).rhint |>
+          regs.(i').rreg |>
+          regs.(i').rhint in
+        regs.(i).rreg <- alloc h i
+    end;
 
+  done;
+  (regs, !spill)
 
 
 
@@ -185,12 +272,14 @@ let parse src =
   ) src;
   p
 
-let t_fact = parse
+let t_pow = parse
   [ "k0:  con 0"
   ; "ni:  con 1234"
   ; "k1:  con 1"
   ; "n0:  phi [ jmp n1 ] [ k1 ni ] ."
+  ; "f1:  phi [ jmp f2 ] [ k1 k1 ] ."
   ; "n1:  sub n0 k1"
+  ; "f2:  add f1 f1"
   ; "jmp: brz n1 end n0"
   ; "end:"
   ]
@@ -210,8 +299,8 @@ let t_fact = parse
   |  |   |
   b7 b6--+
 
-  A simple implementation (that work for non-
-  irreducible control flows) proceeds
+  A simple implementation (that works for
+  non-irreducible control flows) proceeds
   backwards, it would successfully make r1
   live in b2 and b3 but r0 would fail to be
   live in b2. It would become live for the
@@ -235,10 +324,45 @@ let t_irred = parse
   ]
 
 let _ =
+  let p = t_pow in
   let open Printf in
-  let s = liveness t_irred in
-  for i = 0 to Array.length s-1 do
+
+  printf "\n** Liveness analysis:\n";
+  let l = liveness p in
+  for i = 0 to Array.length p -1 do
     printf "%04d:" i;
-    ISet.iter (printf " %04d") s.(i);
+    ISet.iter (printf " %04d") l.(i);
     printf "\n";
+  done;
+
+  printf "\n** Register allocation:\n";
+  let regs = [| "rax"; "rbx"; "rcx" |] in
+  let r, s = regalloc (Array.length regs) p l in
+  if s <> 0 then printf "!! Needs spills !!\n";
+  let bop_str = function
+    | Add -> "add" | Sub -> "sub"
+    | Le -> "cle" | Ge -> "cge" | Lt -> "clt" | Gt -> "cgt" | Eq -> "ceq" | Ne -> "cne" in
+  for i = 0 to Array.length p -1 do
+    let reg i =
+      if r.(i).rreg = -1 then sprintf "%03d" i else regs.(r.(i).rreg) in
+    if r.(i).rreg = -1
+    then printf "%03d:  " i
+    else printf "%s = " (reg i);
+    begin match p.(i) with
+    | ICon k -> printf "%d" k
+    | INop -> ()
+    | IUop (Not, i') -> printf "not %s" (reg i')
+    | IBop (i1, o, i2) ->
+      printf "%s %s %s" (bop_str o) (reg i1) (reg i2)
+    | IBrz (i1, i2, i3) ->
+      printf "brz %s %03d %03d" (reg i1) i2 i3
+    | IJmp i' ->
+      printf "jmp %s" (reg i')
+    | IPhi l ->
+      printf "phi ";
+      List.iter (fun {pjmp; pvar} ->
+        printf "[ %d %s ] " pjmp (reg pvar)
+      ) l
+    end;
+    printf "\n"
   done