Merge pull request #59 from jamesmcnamara/reg-alloc

start of register allocation

fixtures/semant/six.tig

@@ -1 +1 @@
-(1; 2; "foo")
+(1=2; 2=3; "foo")

out.asm

@@ -0,0 +1,40 @@
+.text
+.globl main
+L10:
+sub $sp, $sp, ~60
+L13:
+lw $a0, 0($fp)
+addi $t0, $fp, ~4
+sw $v0, 0($t0)
+addi $t0, $fp, ~8
+sw $s0, 0($t0)
+addi $t0, $fp, ~12
+sw $s1, 0($t0)
+addi $t0, $fp, ~16
+sw $s2, 0($t0)
+addi $t0, $fp, ~20
+sw $s3, 0($t0)
+addi $t0, $fp, ~24
+sw $s4, 0($t0)
+addi $t0, $fp, ~28
+sw $s5, 0($t0)
+addi $t0, $fp, ~32
+sw $s6, 0($t0)
+addi $t0, $fp, ~36
+sw $s7, 0($t0)
+L11move $v0, $t0
+lw $v0, ~4($fp)
+lw $s0, ~8($fp)
+lw $s1, ~12($fp)
+lw $s2, ~16($fp)
+lw $s3, ~20($fp)
+lw $s4, ~24($fp)
+lw $s5, ~28($fp)
+lw $s6, ~32($fp)
+lw $s7, ~36($fp)
+j L12 
+L12:
+addi $sp, $sp, ~60
+jr $ra
+.data
+L11: .ascii "string"

sources.cm

@@ -4,6 +4,8 @@ $/basis.cm
 $/smlnj-lib.cm
 $/ml-yacc-lib.cm
 
+stack.sml
+
 src/main.sml
 src/errormsg.sml
 src/test.sml
@@ -43,3 +45,6 @@ src/flow/graph.sig
 src/flow/graph.sml
 src/flow/makegraph.sml
 src/flow/liveness.sml
+
+src/register/alloc.sml
+src/register/color.sml

src/flow/liveness.sml

@@ -13,7 +13,7 @@ sig
 
   val interferenceGraph: Flow.flowgraph * Flow.Graph.node list -> igraph * (Flow.Graph.node -> Temp.temp list)
   exception NodeDoesNotExist
-  (*val show: TextIO.outstream * igraph -> unit*)
+  val show: igraph -> unit
 
 end
 
@@ -128,7 +128,7 @@ struct
           | addEdges(_,[]) = ()
 
         fun buildInterferenceGraph(liveIn,liveOut) =
-          let val defs = map (fn (node) => (node,(Graph.Table.look(def,node)))) (Graph.nodes(control))
+          let val defs = map (fn (node) => (node,(Graph.Table.look(def,node)))) (Graph.nodes(igraph))
           in
             app (fn (n,d) =>
                     (case d of
@@ -155,4 +155,16 @@ struct
                             efficient as it can be*)
     end
 
+  fun show(IGRAPH{graph=igraph,
+          tnode=tnode,
+          gtemp=gtemp,
+          moves=moves}) =
+    let val nodes = Graph.nodes(igraph)
+        fun printNodes [] = print "COMPLETED\n\n"
+          | printNodes(node::nodes) =
+            (print("-> " ^ Graph.nodename(node) ^ "=" ^ Int.toString(List.length(Graph.adj(node))) ^ "\n"); printNodes(nodes))
+    in
+      print "showing nodes\n\n";
+      printNodes nodes
+    end
 end

src/flow/makegraph.sml

@@ -1,6 +1,5 @@
 signature MAKEGRAPH =
 sig
-  (*val instrs2graph: Assem.instr list -> Flow.flowgraph * Flow.Graph.node list*)
   val instrs2graph: Assem.instr list -> Flow.flowgraph * Flow.Graph.node list
 end
 

src/main.sml

@@ -1,45 +1,41 @@
 signature MAIN =
 sig
     val compile : string -> unit
-
-    val generateCFG : (Symbol.symbol * int * Assem.instr list) list -> (Liveness.igraph * (Flow.Graph.node -> Temp.temp list)) list
 end
 
 structure Main : MAIN =
 struct
 exception NotImplemented
 
-
-fun generateCFG([]) = []
-  | generateCFG((s,i,a)::rest) = (Liveness.interferenceGraph(MakeGraph.instrs2graph(a)))::generateCFG(rest)
-
 fun compile filename =
   let
     fun toAsm(Frame.PROC {body, frame}) =
-        let val assem = CodeGen.codegen frame body
+        let
+          (* Canon *)
+          val stms = Canon.linearize body
+          val stms' = Canon.traceSchedule(Canon.basicBlocks stms)
+          val assem = List.concat(map (CodeGen.codegen frame) stms')
+          val assem' = Frame.procEntryExit2(frame,assem)
+          val { prolog, epilog, body=assem''} = Frame.procEntryExit3(frame, assem')
+          val (assem''', alloc) = RegAlloc.alloc(assem'', frame)
+          val replace = (fn(temp) => valOf(Temp.Table.look(alloc,temp)))
+          val replaceAll = map (fn(instr) => Assem.format replace instr) assem'''
+          val replaced = foldr (fn(a,b) => a ^ b) "" replaceAll
         in
-          Liveness.interferenceGraph(MakeGraph.instrs2graph(assem))
+          ("", (prolog ^ replaced ^ epilog))
         end
-      | toAsm(Frame.STRING (_, _)) = raise NotImplemented
+      | toAsm(Frame.STRING (label, value)) =
+        (Frame.string(label,value),"")
 
     val exp = Parse.parse(filename)
     val _ =  FindEscape.findEscape(exp)
     val ir = Semant.transProg(exp)
-    val asm = (List.map (#1 o toAsm) ir)
-    (*val out = foldr (fn (a, s) => (foldr (fn (i, s) => Assem.format Temp.makestring i ^ s) s a)) "" asm*)
+    val asm = map toAsm ir
+    val (finalString,finalProc) = foldr (fn((string,proc),(stringAcc,procAcc)) => (string ^ stringAcc, proc ^ procAcc)) ("","") asm
+    val final = ".text\n.globl main\n" ^ finalProc ^ ".data\n" ^ finalString
+    val out = TextIO.openOut "out.asm"
   in
-    (* Print the IR fragments. *)
-    (map (fn f =>
-      (case f of
-        MipsFrame.PROC {body, frame} =>
-          Printtree.printtree (TextIO.stdOut, body)
-      | MipsFrame.STRING (l, s) => print(s)))
-      ir);
-
-    print "\n";
-
-    (* Print the Assem. *)
-    (*print out*)
-    ()
+    TextIO.output(out, final);
+    TextIO.closeOut(out)
   end
 end

src/register/alloc.sml

@@ -0,0 +1,48 @@
+signature REG_ALLOC =
+sig
+  structure Frame : FRAME
+  type allocation = Color.allocation
+  val alloc : Assem.instr list * Frame.frame -> Assem.instr list * allocation
+end
+
+structure RegAlloc : REG_ALLOC =
+struct
+  structure Frame = MipsFrame
+  type allocation = Color.allocation
+  exception CannotAllocateError
+
+  fun alloc (instrs, frame) =
+    let val cfg = MakeGraph.instrs2graph(instrs)
+        val (interference as Liveness.IGRAPH{graph=g, tnode=t, gtemp=gt, moves=ms}, liveOut) = Liveness.interferenceGraph(cfg)
+        (*val _ = Liveness.show(interference)*)
+        (* TODO: What to do with the spills? *)
+        val (alloc', spills) =
+          Color.color { interference = (Liveness.IGRAPH{graph=g, tnode=t, gtemp=gt, moves=ms}),
+                        initial = Frame.tempMap,
+                        spillCost = (fn node => 0),
+                        registers = Frame.registers }
+    in
+      (instrs, alloc')
+    end
+end
+
+(*
+datatype igraph = IGRAPH of {graph: graph,
+                             tnode: Temp.temp -> node,
+                             gtemp: node -> Temp.temp,
+                             moves: (node * node) list}
+*)
+
+(*
+  General Algorithm
+  1. Create a stack of nodes to color.
+     Find the first node which can be removed(where adj < K)
+     Remove the node and add it to a stack to color
+     Repeat until no nodes are left
+  2. Add nodes back to the graph and decide on color
+*)
+
+(*
+  TODO: This almost completes part 1.
+        It doesnt remove the edges from nodes that are on the stack
+*)

src/register/color.sml

@@ -0,0 +1,105 @@
+signature COLOR =
+sig
+  structure Frame : FRAME
+  type allocation = Frame.register Temp.Table.table
+  val color: {interference: Liveness.igraph,
+              initial: allocation,
+              spillCost: Graph.node -> int,
+              registers: Frame.register list}
+              -> allocation * Temp.temp list
+end
+
+structure Color : COLOR =
+struct
+  structure Frame : FRAME = MipsFrame
+  type allocation = Frame.register Temp.Table.table
+
+  exception CannotAllocateError
+  exception NotEnoughRegisters
+
+  structure RegisterSet = BinarySetFn(struct type ord_key = Frame.register
+                                     val compare = String.compare
+                              end)
+
+  fun color { interference,
+              initial,
+              spillCost,
+              registers } =
+    let fun findNodeToRemove(nodes) =
+          let fun findNodeToRemoveAcc([],empty) =
+                  if empty then NONE else raise CannotAllocateError
+                | findNodeToRemoveAcc(n::rest,_) =
+                  (if List.length(Graph.adj n) < Frame.sizeOfK
+                   then SOME(n) else findNodeToRemoveAcc(rest,false))
+          in
+            findNodeToRemoveAcc(nodes,true)
+          end
+
+        fun removeNode([],n) = []
+          | removeNode(node::nodes,n) =
+            if Graph.eq(node,n) then nodes else node::removeNode(nodes,n)
+
+        fun buildStack([],s) = s
+          | buildStack(node::nodes,s) =
+            case findNodeToRemove(node::nodes) of
+                NONE => raise CannotAllocateError
+              | SOME(n) =>
+                let val reducedNodes = removeNode(node::nodes,n)
+                in
+                  Stack.push(s,n);
+                  buildStack(reducedNodes,s)
+                end
+
+        val nodeStack = case interference of
+          Liveness.IGRAPH{graph=g, tnode=t, gtemp=gt, moves=ms} =>
+            buildStack(Graph.nodes(g),Stack.empty())
+
+        fun enumerateAdjacentColors(adjacentNodes,alloc,gtemp) =
+          let fun buildSet([],accSet) = accSet
+                | buildSet(node::nodes,accSet) =
+                  case Temp.Table.look(alloc,gtemp(node)) of
+                      NONE => buildSet(nodes,accSet)
+                    | SOME(r) => buildSet(nodes,RegisterSet.add(accSet,r))
+          in
+            buildSet(adjacentNodes,RegisterSet.empty)
+          end
+
+          (* return a set of colors*)
+
+        (* Get all adjacent nodes
+           Enumerate the assigned colors from each adjacent node
+           Take the difference of the available colors and the used colors
+           If there is an unused color, assign it, otherwise raise NotEnoughRegisters
+        *)
+        (* Check each adjacent node. Find the lowest available register *)
+        fun lowestAvailableRegister(node,alloc,gt) =
+          let val adj = Graph.adj(node)
+              val adjacentColors = enumerateAdjacentColors(adj,alloc,gt)
+              val allColors = RegisterSet.addList(RegisterSet.empty, Frame.availableRegisters)
+              val availableColors = RegisterSet.difference(allColors, adjacentColors)
+          in
+            if RegisterSet.numItems(availableColors) > 0 then hd(RegisterSet.listItems(availableColors))
+            else raise NotEnoughRegisters
+          end
+
+        fun apply (stack,alloc,gt) =
+          case Stack.pop(stack) of
+              NONE => alloc
+            | SOME(n) =>
+              case Temp.Table.look(alloc,gt(n)) of
+                  SOME(_) => apply(stack,alloc,gt)
+                | NONE =>
+                  let val register = lowestAvailableRegister(n,alloc,gt)
+                      val updatedAllocation = Temp.Table.enter(alloc,gt(n),register)
+
+                  in
+                    apply(stack,updatedAllocation,gt)
+                  end
+
+  in
+    case interference of
+      Liveness.IGRAPH{graph=g, tnode=t, gtemp=gt, moves=ms} =>
+      (apply(nodeStack,initial,gt), [])
+  end
+
+end

src/semant/assem.sml

@@ -23,8 +23,8 @@ structure Assem = struct
             (explode(saytemp(List.nth(dst,ord i - ord #"0"))) @ f rest)
           | f( #"`":: #"j":: i:: rest) =
             (explode(saylab(List.nth(jump,ord i - ord #"0"))) @ f rest)
-          | f( #"`":: #"`":: rest) = #"`" :: f rest
-          | f( #"`":: _ :: rest) = ErrorMsg.impossible "bad Assem format"
+          | f(#"`":: #"`":: rest) = #"`" :: f rest
+          | f(#"`":: _ :: rest) = ErrorMsg.impossible "bad Assem format"
           | f(c :: rest) = (c :: f rest)
           | f nil = nil
         in implode(f(explode assem))

src/semant/frame.sig

@@ -1,7 +1,7 @@
 signature FRAME =
 sig
   type frame
-  type register 
+  type register
   datatype access = InFrame of int | InReg of Temp.temp
   val newFrame: {name: Temp.label, formals: bool list} -> frame
   val formals: frame -> access list
@@ -14,6 +14,9 @@ sig
   val RA: Temp.temp
   val RV: Temp.temp
   val wordSize: int
+  val sizeOfK: int
+  val registers: register list
+  val availableRegisters: register list
 
   (* For tests *)
   val registersAsTemps: int list
@@ -24,4 +27,9 @@ sig
                 | STRING of Temp.label * string
 
   val tempMap : register Temp.Table.table
+
+  val procEntryExit1: frame * Tree.stm -> Tree.stm
+  val procEntryExit2: frame * Assem.instr list -> Assem.instr list
+  val procEntryExit3: frame * Assem.instr list -> {prolog: string, body: Assem.instr list, epilog: string}
+  val string: Tree.label * string -> string
 end