Provide mechanism to save CFG before register allocation

Makefile.common-jst

@@ -69,7 +69,8 @@ boot_targets = \
   $(boot_ocamlmklib) \
   $(boot_ocamldep) \
   $(boot_ocamlobjinfo) \
-  ocamltest/ocamltest.native
+  ocamltest/ocamltest.native \
+  tools/regalloc/regalloc.exe
 
 boot-compiler: _build/_bootinstall
 	RUNTIME_DIR=$(RUNTIME_DIR) $(dune) build $(ws_boot) $(coverage_dune_flags) $(boot_targets)

asmcomp/asmgen.ml

@@ -77,6 +77,9 @@ let should_save_before_emit () =
 let should_save_cfg_before_emit () =
   should_save_ir_after Compiler_pass.Simplify_cfg && not !start_from_emit
 
+let should_save_cfg_before_regalloc () =
+  should_save_ir_before Compiler_pass.Register_allocation
+
 let linear_unit_info =
   { Linear_format.unit = Compilation_unit.dummy; items = [] }
 
@@ -85,6 +88,8 @@ let new_cfg_unit_info () =
 
 let cfg_unit_info = new_cfg_unit_info ()
 
+let cfg_before_regalloc_unit_info = new_cfg_unit_info ()
+
 module Compiler_pass_map = Map.Make (Compiler_pass)
 
 let (pass_to_cfg : Cfg_format.cfg_unit_info Compiler_pass_map.t) =
@@ -96,10 +101,13 @@ let reset () =
   start_from_emit := false;
   Compiler_pass_map.iter
     (fun pass (cfg_unit_info : Cfg_format.cfg_unit_info) ->
-      if should_save_ir_after pass
+      if should_save_ir_after pass || should_save_ir_before pass
       then (
         cfg_unit_info.unit <- Compilation_unit.get_current_or_dummy ();
-        cfg_unit_info.items <- []))
+        cfg_unit_info.items <- [];
+        cfg_before_regalloc_unit_info.unit
+          <- Compilation_unit.get_current_or_dummy ();
+        cfg_before_regalloc_unit_info.items <- []))
     pass_to_cfg;
   if should_save_before_emit ()
   then (
@@ -134,6 +142,24 @@ let save_cfg f =
   then cfg_unit_info.items <- Cfg_format.(Cfg f) :: cfg_unit_info.items;
   f
 
+let save_cfg_before_regalloc (cfg_with_infos : Cfg_with_infos.t) =
+  (if should_save_cfg_before_regalloc ()
+  then
+    (* CFGs and registers are mutable, so make sure what we will save is a
+       snapshot of the current state. *)
+    let copy x = Marshal.from_string (Marshal.to_string x []) 0 in
+    cfg_before_regalloc_unit_info.items
+      <- Cfg_format.(
+           Cfg_before_regalloc
+             { cfg_with_layout =
+                 copy (Cfg_with_infos.cfg_with_layout cfg_with_infos);
+               cmm_label = Cmm.cur_label ();
+               reg_stamp = Reg.For_testing.get_stamp ();
+               relocatable_regs = copy @@ Reg.all_relocatable_regs ()
+             })
+         :: cfg_before_regalloc_unit_info.items);
+  cfg_with_infos
+
 let write_ir prefix =
   Compiler_pass_map.iter
     (fun pass (cfg_unit_info : Cfg_format.cfg_unit_info) ->
@@ -152,7 +178,15 @@ let write_ir prefix =
   then (
     let filename = Compiler_pass.(to_output_filename Simplify_cfg ~prefix) in
     cfg_unit_info.items <- List.rev cfg_unit_info.items;
-    Cfg_format.save filename cfg_unit_info)
+    Cfg_format.save filename cfg_unit_info);
+  if should_save_cfg_before_regalloc ()
+  then (
+    let filename =
+      Compiler_pass.(to_output_filename Register_allocation ~prefix)
+    in
+    cfg_before_regalloc_unit_info.items
+      <- List.rev cfg_before_regalloc_unit_info.items;
+    Cfg_format.save filename cfg_before_regalloc_unit_info)
 
 let should_emit () = not (should_stop_after Compiler_pass.Linearization)
 
@@ -318,30 +352,27 @@ let available_regs ~stack_slots ~f x =
 let compile_cfg ppf_dump ~funcnames fd_cmm cfg_with_layout =
   let register_allocator = register_allocator fd_cmm in
   let module CSE = Cfg_cse.Cse_generic (CSE) in
-  let cfg_with_infos =
-    cfg_with_layout
-    ++ (fun cfg_with_layout ->
-         match should_vectorize () with
-         | false -> cfg_with_layout
-         | true ->
-           cfg_with_layout
-           ++ cfg_with_layout_profile ~accumulate:true "vectorize"
-                (Vectorize.cfg ppf_dump)
-           ++ pass_dump_cfg_if ppf_dump Flambda_backend_flags.dump_cfg
-                "After vectorize")
-    ++ cfg_with_layout_profile ~accumulate:true "cfg_polling"
-         (Cfg_polling.instrument_fundecl ~future_funcnames:funcnames)
-    ++ cfg_with_layout_profile ~accumulate:true "cfg_zero_alloc_checker"
-         (Zero_alloc_checker.cfg ~future_funcnames:funcnames ppf_dump)
-    ++ cfg_with_layout_profile ~accumulate:true "cfg_comballoc"
-         Cfg_comballoc.run
-    ++ Compiler_hooks.execute_and_pipe Compiler_hooks.Cfg_combine
-    ++ cfg_with_layout_profile ~accumulate:true "cfg_cse" CSE.cfg_with_layout
-    ++ Compiler_hooks.execute_and_pipe Compiler_hooks.Cfg_cse
-    ++ Cfg_with_infos.make
-    ++ cfg_with_infos_profile ~accumulate:true "cfg_deadcode" Cfg_deadcode.run
-  in
-  cfg_with_infos
+  cfg_with_layout
+  ++ (fun cfg_with_layout ->
+       match should_vectorize () with
+       | false -> cfg_with_layout
+       | true ->
+         cfg_with_layout
+         ++ cfg_with_layout_profile ~accumulate:true "vectorize"
+              (Vectorize.cfg ppf_dump)
+         ++ pass_dump_cfg_if ppf_dump Flambda_backend_flags.dump_cfg
+              "After vectorize")
+  ++ cfg_with_layout_profile ~accumulate:true "cfg_polling"
+       (Cfg_polling.instrument_fundecl ~future_funcnames:funcnames)
+  ++ cfg_with_layout_profile ~accumulate:true "cfg_zero_alloc_checker"
+       (Zero_alloc_checker.cfg ~future_funcnames:funcnames ppf_dump)
+  ++ cfg_with_layout_profile ~accumulate:true "cfg_comballoc" Cfg_comballoc.run
+  ++ Compiler_hooks.execute_and_pipe Compiler_hooks.Cfg_combine
+  ++ cfg_with_layout_profile ~accumulate:true "cfg_cse" CSE.cfg_with_layout
+  ++ Compiler_hooks.execute_and_pipe Compiler_hooks.Cfg_cse
+  ++ Cfg_with_infos.make
+  ++ cfg_with_infos_profile ~accumulate:true "cfg_deadcode" Cfg_deadcode.run
+  ++ save_cfg_before_regalloc
   ++ Profile.record ~accumulate:true "regalloc" (fun cfg_with_infos ->
          let cfg_description =
            Regalloc_validate.Description.create

backend/reg.ml

@@ -59,6 +59,13 @@ let dummy =
 let currstamp = ref 0
 let all_relocatable_regs = ref ([] : t list)
 
+module For_testing = struct
+  let get_stamp () = !currstamp
+  let set_state ~stamp ~relocatable_regs =
+    currstamp := stamp;
+    all_relocatable_regs := relocatable_regs  
+end
+
 let create_gen ~name ~typ ~loc =
   let preassigned =
     match loc with
@@ -245,4 +252,4 @@ let same_loc left right =
 let same_loc_fatal_on_unknown ~fatal_message left right =
   match left.loc with
   | Unknown -> Misc.fatal_error fatal_message
-  | Reg _ | Stack _ -> same_loc left right
+  | Reg _ | Stack _ -> same_loc left right

backend/reg.mli

@@ -113,3 +113,8 @@ val same : t -> t -> bool
 val compare : t -> t -> int
 val same_loc : t -> t -> bool
 val same_loc_fatal_on_unknown : fatal_message:string -> t -> t -> bool
+
+module For_testing : sig
+  val get_stamp : unit -> int
+  val set_state : stamp:int -> relocatable_regs:t list -> unit
+end

driver/compenv.ml

@@ -217,6 +217,13 @@ let parse_warnings error v =
 let read_one_param ppf position name v =
   let set name options s =  setter ppf (fun b -> b) name options s in
   let clear name options s = setter ppf (fun b -> not b) name options s in
+  let save_ir ~filter ~setter =
+    if !native_code then begin
+      match decode_compiler_pass ppf v ~name ~filter with
+      | None -> ()
+      | Some pass -> setter pass true
+    end
+  in
   let compat name s =
     let error_if_unset = function
       | true -> true
@@ -476,12 +483,13 @@ let read_one_param ppf position name v =
     set_compiler_pass ppf v ~name Clflags.stop_after ~filter:(fun _ -> true)
 
   | "save-ir-after" ->
-    if !native_code then begin
-      let filter = Clflags.Compiler_pass.can_save_ir_after in
-      match decode_compiler_pass ppf v ~name ~filter with
-      | None -> ()
-      | Some pass -> set_save_ir_after pass true
-    end
+    save_ir
+      ~filter:Clflags.Compiler_pass.can_save_ir_after
+      ~setter:set_save_ir_after
+  | "save-ir-before" ->
+    save_ir
+      ~filter:Clflags.Compiler_pass.can_save_ir_before
+      ~setter:set_save_ir_before
   | "dump-into-file" -> Clflags.dump_into_file := true
   | "dump-into-csv" -> Clflags.dump_into_csv := true
   | "dump-dir" -> Clflags.dump_dir := Some v

driver/main_args.ml

@@ -128,7 +128,17 @@ let mk_save_ir_after ~native f =
                              ~native)
   in
   "-save-ir-after", Arg.Symbol (pass_names, f),
-  " Save intermediate representation after the given compilation pass\
+  " Save intermediate representation after the given compilation pass \
+    (may be specified more than once)."
+
+let mk_save_ir_before ~native f =
+  let pass_names =
+    Clflags.Compiler_pass.(available_pass_names
+                             ~filter:can_save_ir_before
+                             ~native)
+  in
+  "-save-ir-before", Arg.Symbol (pass_names, f),
+  " Save intermediate representation before the given compilation pass \
     (may be specified more than once)."
 
 let mk_dtypes f =
@@ -1141,6 +1151,7 @@ module type Optcomp_options = sig
   val _afl_inst_ratio : int -> unit
   val _function_sections : unit -> unit
   val _save_ir_after : string -> unit
+  val _save_ir_before : string -> unit
   val _probes : unit -> unit
   val _no_probes : unit -> unit
 end;;
@@ -1431,6 +1442,7 @@ struct
     mk_function_sections F._function_sections;
     mk_stop_after ~native:true F._stop_after;
     mk_save_ir_after ~native:true F._save_ir_after;
+    mk_save_ir_before ~native:true F._save_ir_before;
     mk_probes F._probes;
     mk_no_probes F._no_probes;
     mk_i F._i;
@@ -2028,6 +2040,14 @@ module Default = struct
         | None -> () (* this should not occur as we use Arg.Symbol *)
         | Some pass ->
           set_save_ir_after pass true
+
+    let _save_ir_before pass =
+      let module P = Compiler_pass in
+        match P.of_string pass with
+        | None -> () (* this should not occur as we use Arg.Symbol *)
+        | Some pass ->
+          set_save_ir_before pass true
+
     let _thread = set use_threads
     let _verbose = set verbose
     let _version () = Compenv.print_version_string ()

driver/main_args.mli

@@ -252,6 +252,7 @@ module type Optcomp_options = sig
   val _afl_inst_ratio : int -> unit
   val _function_sections : unit -> unit
   val _save_ir_after : string -> unit
+  val _save_ir_before : string -> unit
   val _probes : unit -> unit
   val _no_probes : unit -> unit
 end

driver/maindriver.ml

@@ -69,7 +69,8 @@ let main argv ppf =
            are  incompatible with -pack, -a, -output-obj"
           (String.concat "|"
              (P.available_pass_names ~filter:(fun _ -> true) ~native:false))
-      | Some (P.Middle_end | P.Linearization | P.Simplify_cfg | P.Emit | P.Selection) ->
+      | Some (P.Middle_end | P.Linearization | P.Simplify_cfg | P.Emit
+             | P.Selection | P.Register_allocation) ->
         assert false (* native only *)
     end;
     if !make_archive then begin

driver/optmaindriver.ml

@@ -91,7 +91,8 @@ let main unix argv ppf ~flambda2 =
           Compenv.fatal "Please specify at most one of -pack, -a, -shared, -c, \
                          -output-obj, -instantiate";
       | Some ((P.Parsing | P.Typing | P.Lambda | P.Middle_end | P.Linearization
-              | P.Simplify_cfg | P.Emit | P.Selection) as p) ->
+              | P.Simplify_cfg | P.Emit | P.Selection
+              | P.Register_allocation) as p) ->
         assert (P.is_compilation_pass p);
         Printf.ksprintf Compenv.fatal
           "Options -i and -stop-after (%s) \

file_formats/cfg_format.ml

@@ -19,6 +19,12 @@
 type cfg_item_info =
   | Cfg of Cfg_with_layout.t
   | Data of Cmm.data_item list
+  | Cfg_before_regalloc of {
+    cfg_with_layout : Cfg_with_layout.t;
+    cmm_label: Label.t;
+    reg_stamp: int;
+    relocatable_regs : Reg.t list;
+  } 
 
 type cfg_unit_info =
   {

file_formats/cfg_format.mli

@@ -22,6 +22,12 @@
 type cfg_item_info =
   | Cfg of Cfg_with_layout.t
   | Data of Cmm.data_item list
+  | Cfg_before_regalloc of {
+    cfg_with_layout : Cfg_with_layout.t;
+    cmm_label: Label.t;
+    reg_stamp: int;
+    relocatable_regs : Reg.t list;
+  }
 
 type cfg_unit_info =
   {

tools/regalloc/.ocamlformat-enable

@@ -0,0 +1 @@
+*.ml

tools/regalloc/dune

@@ -0,0 +1,6 @@
+(executable
+ (name regalloc)
+ (modes native)
+ (flags
+  (:standard -w +a-30-40-41-42))
+ (libraries ocamloptcomp))

tools/regalloc/regalloc.ml

@@ -0,0 +1,78 @@
+(* This program simply runs a register allocator on the all the functions saved
+   in a .cmir-cfg-regalloc file. *)
+
+module List = ListLabels
+
+type register_allocator =
+  | GI
+  | IRC
+  | LS
+
+let allocators = ["gi", GI; "irc", IRC; "ls", LS]
+
+external time_include_children : bool -> float
+  = "caml_sys_time_include_children"
+
+let cpu_time () = time_include_children false
+
+let[@ocamlformat "disable"] process_function
+    (register_allocator : register_allocator)
+    (cfg_with_layout : Cfg_with_layout.t) (cmm_label : Label.t)
+    (reg_stamp : int) (relocatable_regs : Reg.t list) =
+  Printf.eprintf "  processing function %S...\n%!"
+    (Cfg_with_layout.cfg cfg_with_layout).fun_name;
+  Printf.eprintf "    %d register(s)...\n%!" (List.length relocatable_regs);
+  Cmm.reset ();
+  Cmm.set_label cmm_label;
+  Reg.For_testing.set_state ~stamp:reg_stamp ~relocatable_regs;
+  let cfg_with_infos = Cfg_with_infos.make cfg_with_layout in
+  let start_time = cpu_time () in
+  let (_ : Cfg_with_infos.t) =
+    match register_allocator with
+    | GI -> Regalloc_gi.run cfg_with_infos
+    | IRC -> Regalloc_irc.run cfg_with_infos
+    | LS -> Regalloc_ls.run cfg_with_infos
+  in
+  let end_time = cpu_time () in
+  Printf.eprintf "  register allocation took %gs...\n%!"
+    (end_time -. start_time);
+  ()
+
+let process_file (file : string) (register_allocator : register_allocator) =
+  Printf.eprintf "processing file %S...\n%!" file;
+  let unit_info, _digest = Cfg_format.restore file in
+  List.iter unit_info.items ~f:(fun (item : Cfg_format.cfg_item_info) ->
+      begin
+        match item with
+        | Cfg _ -> ()
+        | Data _ -> ()
+        | Cfg_before_regalloc
+            { cfg_with_layout; cmm_label; reg_stamp; relocatable_regs } ->
+          process_function register_allocator cfg_with_layout cmm_label
+            reg_stamp relocatable_regs
+      end)
+
+let () =
+  let register_allocator = ref None in
+  let set_register_allocator str =
+    match List.assoc_opt str allocators with
+    | None -> assert false
+    | Some allocator -> register_allocator := Some allocator
+  in
+  let files = ref [] in
+  let args : (Arg.key * Arg.spec * Arg.doc) list =
+    [ ( "-regalloc",
+        Arg.Symbol (List.map allocators ~f:fst, set_register_allocator),
+        "  Choose register allocator" ) ]
+  in
+  let anonymous file = files := file :: !files in
+  Arg.parse args anonymous
+    "run register allocation on a .cmir-cfg-regalloc file";
+  match !register_allocator with
+  | None ->
+    Printf.eprintf
+      "*** error: register allocator was not set (use -regalloc)\n%!";
+    exit 1
+  | Some register_allocator ->
+    List.iter (List.rev !files) ~f:(fun file ->
+        process_file file register_allocator)