Bugfix for register spilling

Author: AdUhTkJm

src/label.ml

@@ -19,66 +19,28 @@ module Label = struct
     }
 
   include struct
-    let _ = fun (_ : t) -> ()
-
-    let sexp_of_t =
-      (fun { name = name__002_; stamp = stamp__004_ } ->
-         let bnds__001_ = ([] : _ Stdlib.List.t) in
-         let bnds__001_ =
-           let arg__005_ = Moon_sexp_conv.sexp_of_int stamp__004_ in
-           (S.List [ S.Atom "stamp"; arg__005_ ] :: bnds__001_ : _ Stdlib.List.t)
-         in
-         let bnds__001_ =
-           let arg__003_ = Moon_sexp_conv.sexp_of_string name__002_ in
-           (S.List [ S.Atom "name"; arg__003_ ] :: bnds__001_ : _ Stdlib.List.t)
-         in
-         S.List bnds__001_
-       : t -> S.t)
-    ;;
-
-    let _ = sexp_of_t
-
-    let equal =
-      (fun a__006_ b__007_ ->
-         if Stdlib.( == ) a__006_ b__007_
-         then true
-         else Stdlib.( = ) (a__006_.stamp : int) b__007_.stamp
-       : t -> t -> bool)
-    ;;
-
-    let _ = equal
-
-    let (hash_fold_t : Ppx_base.state -> t -> Ppx_base.state) =
-      fun hsv arg ->
-      let hsv =
-        let hsv = hsv in
-        hsv
-      in
-      Ppx_base.hash_fold_int hsv arg.stamp
-    ;;
-
-    let _ = hash_fold_t
-
-    let (hash : t -> Ppx_base.hash_value) =
-      let func arg =
-        Ppx_base.get_hash_value
-          (let hsv = Ppx_base.create () in
-           hash_fold_t hsv arg)
-      in
-      fun x -> func x
-    ;;
-
-    let _ = hash
-
-    let compare =
-      (fun a__008_ b__009_ ->
-         if Stdlib.( == ) a__008_ b__009_
-         then 0
-         else Stdlib.compare (a__008_.stamp : int) b__009_.stamp
-       : t -> t -> int)
-    ;;
-
-    let _ = compare
+    let sexp_of_t { name; stamp } =
+      S.List [
+        S.List [ S.Atom "name"; Moon_sexp_conv.sexp_of_string name ];
+        S.List [ S.Atom "stamp"; Moon_sexp_conv.sexp_of_int stamp ]]
+
+    let equal a b =
+         if a == b then true
+         else a.name = b.name && a.stamp = b.stamp
+
+    let hash_fold_t hsv arg =
+      let hsv = 
+       Ppx_base.hash_fold_string hsv arg.name
+      in Ppx_base.hash_fold_int hsv arg.stamp
+
+    let hash arg =
+      Ppx_base.get_hash_value (hash_fold_t (Ppx_base.create ()) arg)
+
+    let compare a b =
+        if a == b then 0
+        else if a.name <> b.name then
+          Stdlib.compare a.name b.name
+        else Stdlib.compare a.stamp b.stamp
   end
 end
 

src/riscv_generate.ml

@@ -868,9 +868,9 @@ let rec do_convert tac (expr: Mcore.expr) =
       let old_vars = !loop_vars in
 
       (* Get the labels *)
-      let loop = Printf.sprintf "%s_%d" label.name label.stamp in
-      let before = Printf.sprintf "before_%s" loop in
-      let exit = Printf.sprintf "exit_%s" loop in
+      let loop = Printf.sprintf "loophead_%s_%d" label.name label.stamp in
+      let before = Printf.sprintf "loopbefore_%s" loop in
+      let exit = Printf.sprintf "loopexit_%s" loop in
 
       Vec.push tac (Jump before);
 
@@ -910,7 +910,7 @@ let rec do_convert tac (expr: Mcore.expr) =
       List.iter2 (fun rd rs -> Vec.push tac (Assign { rd; rs })) !loop_vars results;
 
       (* Jump back to the beginning of the loop. *)
-      let loop_name = Printf.sprintf "%s_%d" label.name label.stamp in 
+      let loop_name = Printf.sprintf "loophead_%s_%d" label.name label.stamp in 
       Vec.push tac (Jump loop_name);
       unit
 

src/riscv_opt.ml

@@ -31,7 +31,10 @@ let exit_fn = Hashtbl.create 256
 let (params: (string, var list) Hashtbl.t) = Hashtbl.create 256
 
 (** Get the basic block with label `name`. *)
-let block_of name = Hashtbl.find basic_blocks name
+let block_of name =
+  match Hashtbl.find_opt basic_blocks name with
+  | None -> failwith (Printf.sprintf "riscv_opt.ml: unknown basic block: %s" name)
+  | Some x -> x
 
 (** Get the body of a basic block. *)
 let body_of name = (block_of name).body |> Vec.to_list

src/riscv_reg.ml

@@ -176,6 +176,7 @@ module Imm = struct
   let to_string imm =
     match imm with
     | IntImm i -> string_of_int i
+    | Int64Imm i -> Int64.to_string i
     | FloatImm f -> string_of_float f
   ;;
 end

src/riscv_reg_spill.ml

@@ -40,6 +40,11 @@ let rpo : RPO.t ref = ref RPO.empty
 
 (* Spill environment: used to build the basic environment, etc. *)
 module SpillEnv = struct
+  (**
+  Data structure for data-flow analysis.
+  Here W stands for working registers (those aren't spilled);
+  S stands for spilled registers.
+  *)
   type spill_info =
     { entryW : SlotSet.t
     ; exitW : SlotSet.t
@@ -145,7 +150,7 @@ let compute_every_inst_nextuse (bl : VBlockLabel.t) : int SlotMap.t Vec.t =
   let block = VProg.get_block !vprog bl in
   let b_liveinfo = Liveness.get_liveinfo !live_info bl in
   let n = Vec.length block.body in
-  let nextUse = Vec.make ~dummy:SlotMap.empty (n + 1) in
+  let nextUse = Vec.of_list (List.init (n + 1) (fun _ -> SlotMap.empty)) in
 
   (* 1. Initialization *)
   SlotMap.iter b_liveinfo.exitNextUse (fun var dist ->
@@ -439,7 +444,7 @@ let apply_min_algorithm (bl : VBlockLabel.t) (nextUse : int SlotMap.t Vec.t) =
 
   (*3. Reload/Spill to be inserted before each instruction, including before Term, so it is n+1 *)
   let body_size = Vec.length block.body in
-  let addInsts = Vec.make ~dummy:(Vec.empty ()) (body_size + 1) in
+  let addInsts = Vec.of_list (List.init (body_size + 1) (fun _ -> Vec.empty ())) in
 
   (* 4. Common apply_inner function, but for clarity, the function to adjust k is passed in *)
   let apply_inner
@@ -472,7 +477,8 @@ let apply_min_algorithm (bl : VBlockLabel.t) (nextUse : int SlotMap.t Vec.t) =
     SlotSet.iter dests (fun var -> w := SlotSet.add !w var);
     let protected = dests in
     (* At this point, protected protects the registers being defined *)
-    let _ = limit_func nextUse w s spill protected (i + 1) adjust_k in
+    if i <> body_size then
+      limit_func nextUse w s spill protected (i + 1) adjust_k;
 
     (* e. Insert reload/spill instructions *)
     SlotSet.iter !reload (fun var -> Vec.push addInsts.![i] (Inst.generate_reload var));
@@ -580,9 +586,10 @@ let spill_reload_func (f_label : VFuncLabel.t) (func : VFunc.t) =
 ;;
 
 (* Main function: used to handle the entire program *)
-let spill_regs (vprog_in : VProg.t) (rpo : RPO.t) =
+let spill_regs (vprog_in : VProg.t) (rpo_arg : RPO.t) =
   vprog := vprog_in;
-  live_info := Liveness.liveness_analysis !vprog rpo;
+  live_info := Liveness.liveness_analysis !vprog rpo_arg;
+  rpo := rpo_arg;
   VFuncMap.iter !vprog.funcs spill_reload_func;
   ()
 ;;

src/riscv_reg_util.ml

@@ -57,7 +57,8 @@ module RPO = struct
   let get_func_rpo (funn : VFuncLabel.t) (rpo : t) : VBlockLabel.t list =
     match VFuncMap.find_opt rpo funn with
     | Some x -> x
-    | None -> failwith "RPO.get_func_rpo: function not found"
+    | None -> failwith
+      (Printf.sprintf "riscv_reg_util.ml: RPO not found for label %s_%d" funn.name funn.stamp)
   ;;
 
   let empty : t = VFuncMap.empty

src/riscv_virtasm_generate.ml

@@ -358,7 +358,14 @@ let gen_fn (f: fn) =
     Hashtbl.add vblocks x {
       body; term = !term; preds =
         Vec.to_list block.pred
-        |> List.map (fun x -> VBlock.NormalEdge (label_of x))
+        |> List.map (fun pred ->
+          (* The labels are fixed for loops in `riscv_generate.ml`. *)
+          if String.starts_with ~prefix:"loophead_" x
+            && not (String.starts_with ~prefix:"loopbefore_" pred) then
+            VBlock.LoopBackEdge (label_of x)
+          else
+            VBlock.NormalEdge (label_of pred)
+        )
     }
   ) blocks;
   ()
@@ -375,13 +382,13 @@ let virtasm_of_ssa (ssa : Riscv_ssa.t list) =
   | ExtArray arr -> gen_extarr arr
   | _ -> failwith "riscv_virtasm_generate.ml: bad toplevel SSA") ssa;
 
-  let funcs = Label.Map.add_list (
+  let funcs = Label.Map.of_list (
     Vec.to_list vfuncs |> List.map (fun (x: VFunc.t) -> (x.funn, x))
-  ) Label.Map.empty in
+  ) in
 
-  let blocks = Label.Map.add_list (
+  let blocks = Label.Map.of_list (
     Hashtbl.to_seq vblocks |> List.of_seq |> List.map (fun (k, v) -> (label_of k, v))
-  ) Label.Map.empty in
+  ) in
 
 
   let out = Printf.sprintf "%s.vasm" !Driver_config.Linkcore_Opt.output_file in

test.py

@@ -68,7 +68,13 @@ def try_remove(path):
 
     for src in cases:
         print(f"Execute task: {src}")
-
+        
+        # Remove files from last time
+        try_remove(f"build/{target}.s.ir")
+        try_remove(f"build/{target}.s.ssa")
+        try_remove(f"build/{target}.s-no-opt.ssa")
+        try_remove(f"build/{target}.s.vasm")
+        
         # Note build-package is ignorant of target. It builds to a common IR.
         os.system(f"moonc build-package src/{src}/{src}.mbt -is-main -std-path {bundled} -o build/{src}.core")