Function Parameter Spilling (#44)

* Separate namespace and parameter tests

* Semi-working

* The easy solution to the register allocation issue

* More parameter spilling tests

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-15 23:49:52.345644
+> Last updated: 2024-05-16 01:16:02.840605
 
 ```
 $ ./main -h

lib/backend/asm.ml

@@ -50,7 +50,7 @@ module Register = struct
       stack and usually needs special care. *)
   let callee_saved_data_registers = [ RBX; R12; R13; R14; R15 ]
 
-  let data_registers = callee_saved_data_registers @ callee_saved_data_registers
+  let data_registers = callee_saved_data_registers @ caller_saved_data_registers
   let parameter_registers = [ RDI; RSI; RDX; RCX; R8; R9 ]
 end
 

lib/backend/asm_clean.ml

@@ -18,7 +18,7 @@ let apply_clean_rules_pair section i =
     when label = Asm.Label.name_of label2 ->
       delete_pair ();
       true
-  | (Push op1, Pop op2 | Pop op1, Push op2) when op1 = op2 ->
+  | Push op1, Pop op2 when op1 = op2 ->
       delete_pair ();
       true
   | (Add (r1, v1), Sub (r2, v2) | Sub (r1, v1), Add (r2, v2))

lib/backend/asm_emit.ml

@@ -4,18 +4,24 @@ module ParameterPassingContext = struct
     mutable regs : Asm.Register.t list;
   }
 
-  let make () = { pos = 0; regs = Asm.Register.parameter_registers }
+  (* this has different semantics from the regalloc allocation type, which is
+     why i'm separating it. register allocation and parameter passing actually
+     have very similar requirements but they're not exactly the same. *)
+  type alloc =
+    | RegisterParam of Asm.Register.t
+    | SpilledParam of int
+
+  let make () = { pos = 1; regs = Asm.Register.parameter_registers }
 
   let get_next ctx =
     if List.is_empty ctx.regs then (
       let pos = ctx.pos in
       ctx.pos <- ctx.pos + 1;
-      ignore pos;
-      failwith "i think should return spill with negative index")
+      SpilledParam pos)
     else
       let result = List.hd ctx.regs in
       ctx.regs <- List.tl ctx.regs;
-      Regalloc.Register result
+      RegisterParam result
 end
 
 let mangle name =
@@ -37,7 +43,7 @@ let align_offset bytes =
 let emit_var regalloc var =
   match VariableMap.find regalloc var with
   | Regalloc.Register reg -> Asm.Operand.Register reg
-  | Spill i -> Asm.Operand.Deref (RBP, (-var_size * i) - var_size)
+  | Spill i -> Asm.Operand.Deref (RBP, -var_size * i)
 
 let emit_oper regalloc = function
   | Operand.Variable var -> emit_var regalloc var
@@ -68,6 +74,41 @@ let emit_restore_registers text registers =
   Asm.Section.add_all text pop_instructions
 
 let emit_call text regalloc name args return_loc_opt =
+  let open Asm.Instruction in
+  let module ParamCtx = ParameterPassingContext in
+  let param_ctx = ParamCtx.make () in
+  let allocs = List.map (fun arg -> (ParamCtx.get_next param_ctx, arg)) args in
+
+  let regs =
+    List.filter_map
+      (fun (alloc, arg) ->
+        match alloc with
+        | ParamCtx.RegisterParam reg -> Some (reg, arg)
+        | SpilledParam _ -> None)
+      allocs
+  in
+  (* since parameter passing registers are a subset of caller saved registers,
+     we should have no problems moving parameters in. *)
+  let reg_movs =
+    List.map (fun (reg, arg) -> Mov (Register reg, emit_oper regalloc arg)) regs
+  in
+
+  let spills =
+    List.filter_map
+      (fun (alloc, arg) ->
+        match alloc with
+        | ParamCtx.RegisterParam _ -> None
+        | SpilledParam i -> Some (i, arg))
+      allocs
+  in
+  (* need to push in reverse order *)
+  let spill_pushes =
+    List.rev_map (fun (_, arg) -> Push (emit_oper regalloc arg)) spills
+  in
+  let max_spill = List.fold_left (fun acc (i, _) -> max acc i) 0 spills in
+  let spill_size = var_size * max_spill in
+  let offset = align_offset spill_size in
+
   let save_registers =
     List.filter
       (fun reg ->
@@ -76,20 +117,48 @@ let emit_call text regalloc name args return_loc_opt =
         | _ -> true)
       Asm.Register.caller_saved_data_registers
   in
+
   emit_save_registers text save_registers;
-  let param_moves =
-    Util.zip_shortest args Asm.Register.parameter_registers
-    |> List.map (fun (arg, reg) ->
-           Asm.Instruction.Mov (Register reg, emit_oper regalloc arg))
-  in
-  Asm.Section.add_all text param_moves;
-  Asm.Section.add text (Asm.Instruction.Call (Label name));
+  Asm.Section.add text (Sub (Register RSP, Intermediate offset));
+  Asm.Section.add_all text spill_pushes;
+  Asm.Section.add_all text reg_movs;
+  Asm.Section.add text (Call (Label name));
   (match return_loc_opt with
   | Some return_loc -> Asm.Section.add text (Mov (return_loc, Register RAX))
   | None -> ());
+  Asm.Section.add text (Add (Register RSP, Intermediate (offset + spill_size)));
   emit_restore_registers text save_registers
 
-let emit_ir text regalloc = function
+let emit_get_param text regalloc param_ctx var =
+  match ParameterPassingContext.get_next param_ctx with
+  | RegisterParam src ->
+      Asm.Section.add text (Mov (emit_var regalloc var, Register src))
+  | SpilledParam i -> (
+      let rbp_offset = var_size + (var_size * i) in
+      let src = Asm.Operand.Deref (RBP, rbp_offset) in
+      let dest = emit_var regalloc var in
+      match dest with
+      | Deref _ ->
+          (* we'll have to use RAX as a temporary *)
+          Asm.Section.add_all text
+            [
+              Push (Register RAX);
+              Mov (Register RAX, src);
+              Mov (dest, Register RAX);
+              Pop (Register RAX);
+            ]
+      | _ -> Asm.Section.add text (Mov (dest, src)))
+
+let emit_return text regalloc op_opt =
+  (match op_opt with
+  | Some op -> Asm.Section.add text (Mov (Register RAX, emit_oper regalloc op))
+  | None -> ());
+
+  emit_restore_registers text Asm.Register.callee_saved_data_registers;
+  Asm.Section.add_all text
+    [ Mov (Register RSP, Register RBP); Pop (Register RBP); Ret ]
+
+let emit_ir text regalloc param_ctx = function
   | Ir.Assign (var, op) ->
       Asm.Section.add text (Mov (emit_var regalloc var, emit_oper regalloc op))
   | Add (var, op, op2) ->
@@ -115,28 +184,15 @@ let emit_ir text regalloc = function
   | DebugPrint op -> emit_call text regalloc debug_print_symbol [ op ] None
   | Call (var, name, args) ->
       emit_call text regalloc (mangle name) args (Some (emit_var regalloc var))
-  | GetParam var -> (
-      let param_passing = ParameterPassingContext.make () in
-      match ParameterPassingContext.get_next param_passing with
-      | Register dest ->
-          Asm.Section.add text (Mov (emit_var regalloc var, Register dest))
-      | Spill _ -> failwith "todo")
-  | Return op_opt ->
-      Option.map
-        (fun op ->
-          Asm.Section.add text (Mov (Register RAX, emit_oper regalloc op)))
-        op_opt
-      |> ignore;
-      emit_restore_registers text Asm.Register.callee_saved_data_registers;
-      Asm.Section.add_all text
-        [ Mov (Register RSP, Register RBP); Pop (Register RBP); Ret ]
+  | GetParam var -> emit_get_param text regalloc param_ctx var
+  | Return op_opt -> emit_return text regalloc op_opt
 
-let emit_bb text cfg regalloc bb =
+let emit_bb text cfg regalloc param_ctx bb =
   Asm.Section.add text
     (Label
        (Asm.Label.make ~is_global:false ~is_external:false
           (Basic_block.label_for bb)));
-  bb |> Basic_block.to_list |> List.iter (emit_ir text regalloc);
+  bb |> Basic_block.to_list |> List.iter (emit_ir text regalloc param_ctx);
   match Basic_block.condition_of bb with
   | Never | Conditional (Constant 0) -> ()
   | Always | Conditional (Constant _) ->
@@ -166,8 +222,7 @@ let emit_cfg ~text cfg regalloc =
         | _ -> acc)
       regalloc 0
   in
-  (* max_spill starts at zero but we need to start from rbp-8 *)
-  let spill_bytes = var_size * (max_spill + 1) in
+  let spill_bytes = var_size * max_spill in
   let stack_bytes = spill_bytes + align_offset spill_bytes in
 
   let entry = Cfg.entry_to cfg in
@@ -185,5 +240,11 @@ let emit_cfg ~text cfg regalloc =
     ];
   (* restore is done at returns *)
   emit_save_registers text Asm.Register.callee_saved_data_registers;
+
+  (* now that we've set up the stack and saved callee-save registers, we can
+     jump to the entrypoint. *)
   Asm.Section.add text (Jmp (Label (Basic_block.label_for entry)));
-  Cfg.blocks_of cfg |> List.iter (emit_bb text cfg regalloc)
+
+  (* we'll need a parameter passing context so that the GetParam IR can work *)
+  let param_ctx = ParameterPassingContext.make () in
+  Cfg.blocks_of cfg |> List.iter (emit_bb text cfg regalloc param_ctx)

lib/backend/regalloc/regalloc.ml

@@ -75,7 +75,7 @@ let linear_scan intervals ordering registers =
   (* must remain sorted by increasing end point *)
   let active : (Variable.t * interval) BatRefList.t = BatRefList.empty () in
 
-  let cur_loc = ref 0 in
+  let cur_loc = ref 1 in
   let next_spill_loc () =
     let result = !cur_loc in
     cur_loc := !cur_loc + 1;

lib/backend/regalloc/regalloc.mli

@@ -1,7 +1,7 @@
 (** Represents the hardware allocated for a variable. [Register reg] means the
     variable has been allocated to a Asm.Register. [Spill i] means the variable
     is to be spilled to the stack. The location [i] will be unique and count up
-    from zero for an allocation scheme produced by [allocate_for]. *)
+    from *one* for an allocation scheme produced by [allocate_for]. *)
 type allocation =
   | Register of Asm.Register.t
   | Spill of int

lib/ir/ir_sim.ml

@@ -61,7 +61,7 @@ let rec run_cfg simulator cfgs cfg =
                  false
              | Ir.Call (result, name, args) ->
                  let called_cfg = find_cfg_by_name cfgs name in
-                 simulator.args <- List.rev_map eval args :: simulator.args;
+                 simulator.args <- List.map eval args :: simulator.args;
                  run_cfg simulator cfgs called_cfg;
                  simulator.args <- List.tl simulator.args;
                  Context.insert simulator.context

lib/user/driver.ml

@@ -49,7 +49,14 @@ let compile paths flags build_dir_loc =
             ]
             cfg liveliness_analysis;
         let instr_ordering = InstrOrdering.make cfg in
-        let registers = Asm.Register.data_registers in
+
+        (* Don't let the allocator use parameter registers, we'll need those in
+           emission. *)
+        let registers =
+          List.filter
+            (fun reg -> not (List.mem reg Asm.Register.parameter_registers))
+            Asm.Register.data_registers
+        in
         let regalloc =
           Regalloc.allocate_for cfg registers liveliness_analysis instr_ordering
         in

lib/util/util.ml

@@ -58,9 +58,4 @@ let basename =
   in
   String.to_seq >> List.of_seq >> List.rev >> basename_aux
 
-let rec zip_shortest lst1 lst2 =
-  match (lst1, lst2) with
-  | h1 :: t1, h2 :: t2 -> (h1, h2) :: zip_shortest t1 t2
-  | _, _ -> []
-
 let pp_of string_of fmt x = Format.fprintf fmt "%s" (string_of x)

lib/util/util.mli

@@ -42,10 +42,6 @@ val merge_paths : string list -> string
     ['/']. *)
 val basename : string -> string
 
-(** [zip_shortest [a1; a2; ... an] [b1; b2; ... bm]] is
-    [(a1, b1); (a2, b2); ... (ak, bk)] where [k = min(n, m)]. *)
-val zip_shortest : 'a list -> 'b list -> ('a * 'b) list
-
 (** [pp_of string_of] is a pretty printer for a type with the string conversion
     function [string_of] that simply prints the result of [string_of] inline. *)
 val pp_of : ('a -> string) -> Format.formatter -> 'a -> unit

test/e2e/namespace.x86istmb

@@ -0,0 +1,9 @@
+namespace my {
+	func display() {
+		print 3110
+	}
+}
+
+func main() {
+	my::display()
+}

test/e2e/param_spill_simple.x86istmb

@@ -0,0 +1,13 @@
+func one_spill(a: Int, b: Int, c: Int, d: Int, e: Int, f: Int, g: Int) {
+    print a
+    print b
+    print c
+    print d
+    print e
+    print f
+    print g
+}
+
+func main() {
+	one_spill(1, 2, 3, 4, 5, 6, 7)
+}

test/e2e/param_spill_three.x86istmb

@@ -0,0 +1,15 @@
+func three_spill(a: Int, b: Int, c: Int, d: Int, e: Int, f: Int, g: Int, h: Int, i: Int) {
+    print a
+    print b
+    print c
+    print d
+    print e
+    print f
+    print g
+    print h
+    print i
+}
+
+func main() {
+	three_spill(1, 2, 3, 4, 5, 6, 7, 8, 9)
+}

test/e2e/param_spill_with_reg_spill.x86istmb

@@ -0,0 +1,36 @@
+func one_spill(a: Int, b: Int, c: Int, d: Int, e: Int, f: Int, g: Int) {
+    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 p = 15
+    let q = 16
+    let r = 17
+
+    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
+    print r
+}
+
+func main() {
+	one_spill(1, 2, 3, 4, 5, 6, 7)
+}

test/e2e/parameters.x86istmb

@@ -0,0 +1,7 @@
+func display(val: Int) {
+    print val
+}
+
+func main() {
+	display(40)
+}