Tweak condition forcing linscan (register allocation) (ocaml-flambda#3913)

Author: xclerc

asmcomp/asmgen.ml

@@ -156,8 +156,15 @@ let write_ir prefix =
 
 let should_emit () = not (should_stop_after Compiler_pass.Linearization)
 
-let should_use_linscan fun_codegen_options =
-  !use_linscan || List.mem Cmm.Use_linscan_regalloc fun_codegen_options
+(* note: `should_use_linscan` relies on the state of the `Reg` module, as the
+   list of temporaries is retrieved to be compared to the threshold. *)
+let should_use_linscan fd =
+  !use_linscan
+  || List.mem Cmm.Use_linscan_regalloc fd.fun_codegen_options
+  || List.compare_length_with
+       (Reg.all_relocatable_regs ())
+       !Flambda_backend_flags.regalloc_linscan_threshold
+     > 0
 
 let if_emit_do f x = if should_emit () then f x else ()
 
@@ -290,7 +297,7 @@ type register_allocator =
 
 let register_allocator fd : register_allocator =
   match String.lowercase_ascii !Flambda_backend_flags.regalloc with
-  | "" | "cfg" -> if should_use_linscan fd.fun_codegen_options then LS else IRC
+  | "" | "cfg" -> if should_use_linscan fd then LS else IRC
   | "gi" -> GI
   | "irc" -> IRC
   | "ls" -> LS

driver/flambda_backend_args.ml

@@ -42,6 +42,11 @@ let mk_dcfg_invariants f =
 let mk_regalloc f =
   "-regalloc", Arg.String f, " Select the register allocator"
 
+let mk_regalloc_linscan_threshold f =
+  "-regalloc-linscan-threshold",
+  Arg.Int f,
+  (Printf.sprintf " Use linscan on functions with more temporaries than the threshold (default is %d)"Flambda_backend_flags.default_regalloc_linscan_threshold)
+
 let mk_regalloc_param f =
   "-regalloc-param", Arg.String f, " Pass a parameter to the register allocator"
 
@@ -725,6 +730,7 @@ module type Flambda_backend_options = sig
   val dcfg : unit -> unit
   val dcfg_invariants : unit -> unit
   val regalloc : string -> unit
+  val regalloc_linscan_threshold : int -> unit
   val regalloc_param : string -> unit
   val regalloc_validate : unit -> unit
   val no_regalloc_validate : unit -> unit
@@ -860,6 +866,7 @@ struct
     mk_dcfg F.dcfg;
     mk_dcfg_invariants F.dcfg_invariants;
     mk_regalloc F.regalloc;
+    mk_regalloc_linscan_threshold F.regalloc_linscan_threshold;
     mk_regalloc_param F.regalloc_param;
     mk_regalloc_validate F.regalloc_validate;
     mk_no_regalloc_validate F.no_regalloc_validate;
@@ -1026,6 +1033,7 @@ module Flambda_backend_options_impl = struct
   let dcfg = set' Flambda_backend_flags.dump_cfg
   let dcfg_invariants = set' Flambda_backend_flags.cfg_invariants
   let regalloc x = Flambda_backend_flags.regalloc := x
+  let regalloc_linscan_threshold x = Flambda_backend_flags.regalloc_linscan_threshold := x
   let regalloc_param x = Flambda_backend_flags.regalloc_params := x :: !Flambda_backend_flags.regalloc_params
   let regalloc_validate = set' Flambda_backend_flags.regalloc_validate
   let no_regalloc_validate = clear' Flambda_backend_flags.regalloc_validate
@@ -1378,6 +1386,7 @@ module Extra_params = struct
     | "ocamlcfg" -> set' Flambda_backend_flags.use_ocamlcfg
     | "cfg-invariants" -> set' Flambda_backend_flags.cfg_invariants
     | "regalloc" -> set_string Flambda_backend_flags.regalloc
+    | "regalloc-linscan-threshold" -> set_int' Flambda_backend_flags.regalloc_linscan_threshold
     | "regalloc-param" -> add_string Flambda_backend_flags.regalloc_params
     | "regalloc-validate" -> set' Flambda_backend_flags.regalloc_validate
     | "vectorize" -> set' Flambda_backend_flags.vectorize

driver/flambda_backend_args.mli

@@ -29,6 +29,7 @@ module type Flambda_backend_options = sig
   val dcfg : unit -> unit
   val dcfg_invariants : unit -> unit
   val regalloc : string -> unit
+  val regalloc_linscan_threshold : int -> unit
   val regalloc_param : string -> unit
   val regalloc_validate : unit -> unit
   val no_regalloc_validate : unit -> unit

driver/flambda_backend_flags.ml

@@ -17,6 +17,8 @@ let use_ocamlcfg = ref true             (* -[no-]ocamlcfg *)
 let dump_cfg = ref false                (* -dcfg *)
 let cfg_invariants = ref false          (* -dcfg-invariants *)
 let regalloc = ref ""                   (* -regalloc *)
+let default_regalloc_linscan_threshold = 100_000
+let regalloc_linscan_threshold = ref max_int (* -regalloc-linscan-threshold *)
 let regalloc_params = ref ([] : string list)  (* -regalloc-param *)
 let regalloc_validate = ref true        (* -[no-]regalloc-validate *)
 

driver/flambda_backend_flags.mli

@@ -18,6 +18,8 @@ val use_ocamlcfg : bool ref
 val dump_cfg : bool ref
 val cfg_invariants : bool ref
 val regalloc : string ref
+val default_regalloc_linscan_threshold : int
+val regalloc_linscan_threshold : int ref
 val regalloc_params : string list ref
 val regalloc_validate : bool ref