Cold continuations: backend changes (ocaml-flambda#1576)

backend/cfg/cfg.ml

@@ -41,7 +41,8 @@ type basic_block =
     mutable exn : Label.t option;
     mutable can_raise : bool;
     mutable is_trap_handler : bool;
-    mutable dead : bool
+    mutable dead : bool;
+    mutable cold : bool
   }
 
 type t =

backend/cfg/cfg.mli

@@ -54,10 +54,11 @@ type basic_block =
     mutable is_trap_handler : bool;
         (** Is this block a trap handler (i.e. is it an exn successor of another
             block) or not? *)
-    mutable dead : bool
+    mutable dead : bool;
         (** This block must be unreachable from function entry. This field is
             set during cfg construction (if trap stacks are unresolved) and used
             during dead block elimination for checking. *)
+    mutable cold : bool
         (* CR-someday gyorsh: The current implementation allows multiple
            pushtraps in each block means that different trap stacks are
            associated with the block at different points. At most one

backend/cfg/cfg_with_layout.ml

@@ -388,6 +388,18 @@ let reorder_blocks_random ?random_state t =
   in
   set_layout t (DLL.of_list new_layout)
 
+let reorder_blocks ~comparator t =
+  (* CR ncourant: this is only ever called with a boolean comparator, we could
+     do better. Or maybe we should write stable_sort on DLL to avoid the
+     conversions? *)
+  (* Ensure entry block remains first *)
+  let original_layout = DLL.to_list (layout t) in
+  let new_layout =
+    List.hd original_layout
+    :: List.stable_sort comparator (List.tl original_layout)
+  in
+  set_layout t (DLL.of_list new_layout)
+
 let iter_instructions :
     t ->
     instruction:(Cfg.basic Cfg.instruction -> unit) ->

backend/cfg/cfg_with_layout.mli

@@ -103,6 +103,8 @@ val dump : Format.formatter -> t -> msg:string -> unit
     [Random.State.bool]. *)
 val reorder_blocks_random : ?random_state:Random.State.t -> t -> unit
 
+val reorder_blocks : comparator:(Label.t -> Label.t -> int) -> t -> unit
+
 val iter_instructions :
   t ->
   instruction:(Cfg.basic Cfg.instruction -> unit) ->

backend/cfg/cfgize.ml

@@ -439,19 +439,21 @@ let extract_block_info : State.t -> Mach.instruction -> block_info =
    return. *)
 let fallthrough_label : Label.t = -1
 
-(* [add_blocks instr state ~starts_with_pushtrap ~start ~next] adds the block
-   beginning at [instr] with label [start], and all recursively-reachable blocks
-   to [state]. [next] is the label of the block to be executed after the one
-   beginning at [instr]. [starts_with_pushtrap] indicates whether the block
-   should be prefixed with a pushtrap instruction (to the passed label). *)
+(* [add_blocks instr state ~starts_with_pushtrap ~start ~next ~is_cold] adds the
+   block beginning at [instr] with label [start], and all recursively-reachable
+   blocks to [state]. [next] is the label of the block to be executed after the
+   one beginning at [instr]. [starts_with_pushtrap] indicates whether the block
+   should be prefixed with a pushtrap instruction (to the passed label).
+   [is_cold] indicates whether to put the blocks in the cold section. *)
 let rec add_blocks :
     Mach.instruction ->
     State.t ->
     starts_with_pushtrap:Label.t option ->
     start:Label.t ->
     next:Label.t ->
+    is_cold:bool ->
     unit =
- fun instr state ~starts_with_pushtrap ~start ~next ->
+ fun instr state ~starts_with_pushtrap ~start ~next ~is_cold ->
   let { instrs; last; terminator } = extract_block_info state instr in
   let terminate_block ~trap_actions terminator =
     let body = instrs in
@@ -499,7 +501,8 @@ let rec add_blocks :
           can_raise;
           (* See [update_trap_handler_blocks] *)
           is_trap_handler = false;
-          dead = false
+          dead = false;
+          cold = is_cold
         }
   in
   let prepare_next_block () =
@@ -511,6 +514,7 @@ let rec add_blocks :
       let start = Cmm.new_label () in
       let add_next_block () =
         add_blocks last.next state ~starts_with_pushtrap:None ~start ~next
+          ~is_cold
       in
       start, add_next_block
   in
@@ -547,8 +551,10 @@ let rec add_blocks :
         (copy_instruction state last
            ~desc:(terminator_of_test test ~label_false ~label_true));
       let next, add_next_block = prepare_next_block () in
-      add_blocks ifso state ~starts_with_pushtrap:None ~start:label_true ~next;
-      add_blocks ifnot state ~starts_with_pushtrap:None ~start:label_false ~next;
+      add_blocks ifso state ~starts_with_pushtrap:None ~start:label_true ~next
+        ~is_cold;
+      add_blocks ifnot state ~starts_with_pushtrap:None ~start:label_false ~next
+        ~is_cold;
       add_next_block ()
     | Iswitch (indexes, cases) ->
       let case_labels = Array.map (fun _ -> Cmm.new_label ()) cases in
@@ -559,26 +565,28 @@ let rec add_blocks :
       Array.iteri
         (fun idx case ->
           add_blocks case state ~starts_with_pushtrap:None
-            ~start:case_labels.(idx) ~next)
+            ~start:case_labels.(idx) ~next ~is_cold)
         cases;
       add_next_block ()
     | Icatch (_rec, _trap_stack, handlers, body) ->
       let handlers =
         List.map
-          (fun (handler_id, _trap_stack, handler, _is_cold) ->
+          (fun (handler_id, _trap_stack, handler, is_cold) ->
             let handler_label = State.add_catch_handler state ~handler_id in
-            handler_label, handler)
+            handler_label, handler, is_cold)
           handlers
       in
       let body_label = Cmm.new_label () in
       terminate_block ~trap_actions:[]
         (copy_instruction_no_reg state last ~desc:(Cfg.Always body_label));
       let next, add_next_block = prepare_next_block () in
-      add_blocks body state ~starts_with_pushtrap:None ~start:body_label ~next;
+      add_blocks body state ~starts_with_pushtrap:None ~start:body_label ~next
+        ~is_cold;
       List.iter
-        (fun (handler_label, handler) ->
+        (fun (handler_label, handler, is_handler_cold) ->
           add_blocks handler state ~starts_with_pushtrap:None
-            ~start:handler_label ~next)
+            ~start:handler_label ~next
+            ~is_cold:(is_cold || is_handler_cold))
         handlers;
       add_next_block ()
     | Iexit (handler_id, trap_actions) ->
@@ -601,9 +609,10 @@ let rec add_blocks :
       let next, add_next_block = prepare_next_block () in
       State.add_iend_with_poptrap state (get_end body);
       State.add_exception_handler state label_handler;
-      add_blocks body state ~starts_with_pushtrap ~start:label_body ~next;
+      add_blocks body state ~starts_with_pushtrap ~start:label_body ~next
+        ~is_cold;
       add_blocks handler state ~starts_with_pushtrap:None ~start:label_handler
-        ~next;
+        ~next ~is_cold;
       add_next_block ()
     | Iraise raise_kind ->
       terminate_block ~trap_actions:[]
@@ -801,7 +810,8 @@ let fundecl :
         exn = None;
         can_raise = false;
         is_trap_handler = false;
-        dead = false
+        dead = false;
+        cold = false
       };
   State.add_block state ~label:tailrec_label
     ~block:
@@ -815,10 +825,11 @@ let fundecl :
         exn = None;
         can_raise = false;
         is_trap_handler = false;
-        dead = false
+        dead = false;
+        cold = false
       };
   add_blocks fun_body state ~starts_with_pushtrap:None ~start:start_label
-    ~next:fallthrough_label;
+    ~next:fallthrough_label ~is_cold:false;
   update_trap_handler_blocks state cfg;
   (* note: `Stack_offset_and_exn.update_cfg` may add edges to the graph, and
      should hence be executed before
@@ -841,4 +852,10 @@ let fundecl :
       Eliminate_dead_code.run_dead_block cfg_with_layout;
       if simplify_terminators then Simplify_terminator.run cfg)
     ();
+  Cfg_with_layout.reorder_blocks
+    ~comparator:(fun label1 label2 ->
+      let block1 = Cfg.get_block_exn cfg label1 in
+      let block2 = Cfg.get_block_exn cfg label2 in
+      Bool.compare block1.cold block2.cold)
+    cfg_with_layout;
   cfg_with_layout

backend/cfg/linear_to_cfg.ml

@@ -182,7 +182,8 @@ let create_empty_block t start ~stack_offset ~traps =
       stack_offset;
       is_trap_handler = false;
       can_raise = false;
-      dead = false
+      dead = false;
+      cold = false
     }
   in
   record_traps t start traps;

backend/cfg/merge_straightline_blocks.ml

@@ -77,6 +77,7 @@ let rec merge_blocks (removed : Label.Set.t)
             b1_block.terminator <- b2_block.terminator;
             b1_block.exn <- b2_block.exn;
             b1_block.can_raise <- b2_block.can_raise;
+            b1_block.cold <- b1_block.cold || b2_block.cold;
             (* modify b2 *)
             b2_block.predecessors <- Label.Set.empty;
             Label.Set.iter

backend/regalloc/regalloc_utils.ml

@@ -543,7 +543,8 @@ let insert_block :
           exn = None;
           can_raise = false;
           is_trap_handler = false;
-          dead = predecessor_block.dead
+          dead = predecessor_block.dead;
+          cold = predecessor_block.cold
         }
       in
       Cfg_with_layout.add_block cfg_with_layout block

tests/backend/regalloc_validator/check_regalloc_validation.ml

@@ -76,7 +76,8 @@ module Block = struct
       exn;
       can_raise;
       is_trap_handler = false;
-      dead = false
+      dead = false;
+      cold = false
     }
 end
 
@@ -181,6 +182,7 @@ let entry_label =
            predecessors = Label.Set.empty;
            stack_offset = 0;
            dead = false;
+           cold = false;
            terminator =
              { desc = Return;
                arg = [| Proc.phys_reg 0 |];