Move down the stack check if possible (ocaml-flambda#2373)

backend/.ocamlformat-enable

@@ -19,5 +19,7 @@ regalloc/**/*.mli
 amd64/simd*.ml
 amd64/stack_check.ml
 arm64/simd*.ml
+amd64/stack_check.ml
+arm64/stack_check.ml
 generic_fns.ml
 generic_fns.mli

backend/amd64/emit.mlp

@@ -487,7 +487,7 @@ let emit_stack_realloc () =
       cfi_adjust_cfa_offset 8;
       (* measured in words *)
       emit_call (Cmm.global_symbol "caml_call_realloc_stack");
-      I.pop r10; (* ignored *)
+      I.add (int 8) rsp;
       cfi_adjust_cfa_offset (-8);
       I.jmp (label sc_return)
     end

backend/arm64/stack_check.ml

@@ -18,15 +18,14 @@
 
 let stack_threshold_size = 0
 
-let frame_size
-  : stack_offset:int -> frame_required:bool -> num_stack_slots:int array ->  int
-  = fun ~stack_offset:_ ~frame_required:_ ~num_stack_slots:_ ->
+let frame_size :
+    stack_offset:int -> frame_required:bool -> num_stack_slots:int array -> int
+    =
+ fun ~stack_offset:_ ~frame_required:_ ~num_stack_slots:_ ->
   Misc.fatal_error "stack checks are not supported on arm64"
 
-let linear
-  : Linear.fundecl -> Linear.fundecl
-  = fun fundecl ->
+let linear : Linear.fundecl -> Linear.fundecl =
+ fun fundecl ->
   match Config.runtime5 with
   | false -> fundecl
-  | true ->
-    Misc.fatal_error "stack checks are not supported on arm64"
+  | true -> Misc.fatal_error "stack checks are not supported on arm64"

backend/asmgen.ml

@@ -335,6 +335,10 @@ let compile_fundecl ~ppf_dump ~funcnames fd_cmm =
              we would have to recompute it here. Recomputing it here breaks the CI because
              the liveness_analysis algorithm does not work properly after register allocation. *)
         ++ Profile.record ~accumulate:true "peephole_optimize_cfg" Peephole_optimize.peephole_optimize_cfg
+        ++ (fun (cfg_with_layout : Cfg_with_layout.t) ->
+          match !Flambda_backend_flags.cfg_stack_checks with
+          | false -> cfg_with_layout
+          | true -> Cfg_stack_checks.cfg cfg_with_layout)
         ++ Profile.record ~accumulate:true "save_cfg" save_cfg
         ++ Profile.record ~accumulate:true "cfg_reorder_blocks"
              (reorder_blocks_random ppf_dump)
@@ -376,6 +380,10 @@ let compile_fundecl ~ppf_dump ~funcnames fd_cmm =
                 ~simplify_terminators:true)
         ++ Compiler_hooks.execute_and_pipe Compiler_hooks.Cfg
         ++ pass_dump_cfg_if ppf_dump Flambda_backend_flags.dump_cfg "After cfgize"
+        ++ (fun (cfg_with_layout : Cfg_with_layout.t) ->
+          match !Flambda_backend_flags.cfg_stack_checks with
+          | false -> cfg_with_layout
+          | true -> Cfg_stack_checks.cfg cfg_with_layout)
         ++ Profile.record ~accumulate:true "save_cfg" save_cfg
         ++ Profile.record ~accumulate:true "cfg_reorder_blocks"
              (reorder_blocks_random ppf_dump)
@@ -385,7 +393,10 @@ let compile_fundecl ~ppf_dump ~funcnames fd_cmm =
   ++ Profile.record ~accumulate:true "scheduling" Scheduling.fundecl
   ++ pass_dump_linear_if ppf_dump dump_scheduling "After instruction scheduling"
   ++ Profile.record ~accumulate:true "save_linear" save_linear
-  ++ Stack_check.linear
+  ++ (fun (fd : Linear.fundecl) ->
+    match !Flambda_backend_flags.cfg_stack_checks with
+    | false -> Stack_check.linear fd
+    | true -> fd)
   ++ Profile.record ~accumulate:true "emit_fundecl" emit_fundecl
 
 let compile_data dl =

backend/cfg/cfg.ml

@@ -191,10 +191,18 @@ let get_block_exn t label =
 
 let can_raise_interproc block = block.can_raise && Option.is_none block.exn
 
-let first_instruction_id (block : basic_block) : int =
+type 'a instr_mapper = { f : 'b. 'b instruction -> 'a } [@@unboxed]
+
+let map_first_instruction (block : basic_block) (t : 'a instr_mapper) =
   match DLL.hd block.body with
-  | None -> block.terminator.id
-  | Some first_instr -> first_instr.id
+  | None -> t.f block.terminator
+  | Some first_instr -> t.f first_instr
+
+let first_instruction_id (block : basic_block) : int =
+  map_first_instruction block { f = (fun instr -> instr.id) }
+
+let first_instruction_stack_offset (block : basic_block) : int =
+  map_first_instruction block { f = (fun instr -> instr.stack_offset) }
 
 let fun_name t = t.fun_name
 
@@ -556,3 +564,21 @@ let string_of_irc_work_list = function
   | Frozen -> "frozen"
   | Work_list -> "work_list"
   | Active -> "active"
+
+let make_instruction ~desc ?(arg = [||]) ?(res = [||]) ?(dbg = Debuginfo.none)
+    ?(fdo = Fdo_info.none) ?(live = Reg.Set.empty) ~stack_offset ~id
+    ?(irc_work_list = Unknown_list) ?(ls_order = 0) ?(available_before = None)
+    ?(available_across = None) () =
+  { desc;
+    arg;
+    res;
+    dbg;
+    fdo;
+    live;
+    stack_offset;
+    id;
+    irc_work_list;
+    ls_order;
+    available_before;
+    available_across
+  }

backend/cfg/cfg.mli

@@ -123,6 +123,8 @@ val can_raise_interproc : basic_block -> bool
 
 val first_instruction_id : basic_block -> int
 
+val first_instruction_stack_offset : basic_block -> int
+
 val mem_block : t -> Label.t -> bool
 
 val add_block_exn : t -> basic_block -> unit
@@ -185,3 +187,19 @@ val string_of_irc_work_list : irc_work_list -> string
 val dump_basic : Format.formatter -> basic -> unit
 
 val dump_terminator : ?sep:string -> Format.formatter -> terminator -> unit
+
+val make_instruction :
+  desc:'a ->
+  ?arg:Reg.t array ->
+  ?res:Reg.t array ->
+  ?dbg:Debuginfo.t ->
+  ?fdo:Fdo_info.t ->
+  ?live:Reg.Set.t ->
+  stack_offset:int ->
+  id:int ->
+  ?irc_work_list:irc_work_list ->
+  ?ls_order:int ->
+  ?available_before:Reg_availability_set.t option ->
+  ?available_across:Reg_availability_set.t option ->
+  unit ->
+  'a instruction

backend/cfg/cfg_dominators.ml

@@ -21,7 +21,8 @@ type dominator_tree =
   }
 
 type t =
-  { doms : doms;
+  { entry_label : Label.t;
+    doms : doms;
     dominance_frontiers : dominance_frontiers;
     dominator_forest : dominator_tree list
   }
@@ -418,7 +419,7 @@ let build : Cfg.t -> t =
   let doms = compute_doms cfg in
   let dominance_frontiers = compute_dominance_frontiers cfg doms in
   let dominator_forest = compute_dominator_forest cfg doms in
-  { doms; dominance_frontiers; dominator_forest }
+  { entry_label = cfg.entry_label; doms; dominance_frontiers; dominator_forest }
 
 let is_dominating t left right = is_dominating t.doms left right
 
@@ -434,6 +435,18 @@ let find_dominance_frontier t label =
 
 let dominator_forest t = t.dominator_forest
 
+let dominator_tree_for_entry_point t =
+  match
+    List.find_opt t.dominator_forest ~f:(fun tree ->
+        Label.equal tree.label t.entry_label)
+  with
+  | None ->
+    fatal
+      "Cfg_dominators.dominator_tree_for_entry_point: no tree for entry point \
+       (label %d)"
+      t.entry_label
+  | Some tree -> tree
+
 let iter_breadth_dominator_forest t ~f =
   List.iter t.dominator_forest ~f:(fun dominator_tree ->
       iter_breadth_dominator_tree dominator_tree ~f)

backend/cfg/cfg_dominators.mli

@@ -1,6 +1,6 @@
 [@@@ocaml.warning "+a-4-30-40-41-42"]
 
-(* Dominator-related utility functions. *)
+(** Dominator-related utility functions. *)
 
 type dominator_tree = private
   { label : Label.t;
@@ -9,29 +9,35 @@ type dominator_tree = private
 
 type t
 
-val build : Cfg.t -> t
-(* Computes all dominator-related information, in particular immediate
+(** Computes all dominator-related information, in particular immediate
    dominators, dominance frontiers, and dominator forest for the passed CFG. *)
+val build : Cfg.t -> t
 
-val is_dominating : t -> Label.t -> Label.t -> bool
-(* [is_dominating doms x y] is [true] iff [x] is dominating [y] according to
+(** [is_dominating doms x y] is [true] iff [x] is dominating [y] according to
    [doms]. That is, all paths from the entry node to [y] go through [x]. All
    edges, regular and exceptional are treated the same way. *)
+val is_dominating : t -> Label.t -> Label.t -> bool
 
-val is_strictly_dominating : t -> Label.t -> Label.t -> bool
-(* [is_strictly_dominating doms x y] is [true] iff [x] is strictly dominating
+(** [is_strictly_dominating doms x y] is [true] iff [x] is strictly dominating
    [y] according to [doms]. That is, [is_dominating doms x y = true] and [x] is
    not equal [y]. All edges, regular and exceptional are treated the same way.*)
+val is_strictly_dominating : t -> Label.t -> Label.t -> bool
 
-val find_dominance_frontier : t -> Label.t -> Label.Set.t
-(* [find_dominance_frontier doms label] returns the dominance frontier for
+(** [find_dominance_frontier doms label] returns the dominance frontier for
    [label] according to [doms]. The definition we use is the following: "the
    dominance frontier of a node n is the set of all nodes m such that n
    dominates a predecessor of m, but does not strictly dominate m itself". *)
+val find_dominance_frontier : t -> Label.t -> Label.Set.t
 
+(** Returns all dominator trees. Typically one of these will correspond to the
+   entry point of the program and the remainder to dead code. *)
 val dominator_forest : t -> dominator_tree list
 
-val iter_breadth_dominator_forest : t -> f:(Label.t -> unit) -> unit
-(* [iter_breadth_dominator_forest doms ~f] iterates over the dominator forest
+(** Returns the unique dominator tree rooted at the entry point of the program
+   (thus ignoring any isolated trees that correspond to dead code). *)
+val dominator_tree_for_entry_point : t -> dominator_tree
+
+(** [iter_breadth_dominator_forest doms ~f] iterates over the dominator forest
    from [doms] in a breadth-first manner (iterating over a whole tree of the
    forest before moving to the next tree), applying [f] to visited nodes. *)
+val iter_breadth_dominator_forest : t -> f:(Label.t -> unit) -> unit