chapter 20: register allocation with conservative coalescing

bin/main.ml

@@ -206,7 +206,7 @@ let debug =
     let mk_regalloc_opts dump_lvl =
       Settings.
         {
-          spill_info = dump_lvl >= 1;
+          debug_msg = dump_lvl >= 1;
           interference_ncol = dump_lvl >= 2;
           interference_graphviz = dump_lvl >= 3;
           liveness = dump_lvl >= 4;

lib/backend/regalloc.ml

@@ -26,6 +26,7 @@ let get_operands = function
   | Label _ | Call _ | Ret | Cdq _ | JmpCC _ | Jmp _ -> []
   | Pop _ -> failwith "Internal error" [@coverage off]
 
+(* map functon f over all the operands in an instruction *)
 let replace_ops f i =
   match i with
   | Mov (t, src, dst) -> Mov (t, f src, f dst)
@@ -42,7 +43,7 @@ let replace_ops f i =
   | SetCC (code, dst) -> SetCC (code, f dst)
   | Push v -> Push (f v)
   | Label _ | Call _ | Ret | Cdq _ | Jmp _ | JmpCC _ -> i
-  | Pop _ -> failwith "We shouldn't use this yet" [@coverage off]
+  | Pop _ -> failwith "Shouldn't use this yet" [@coverage off]
 
 let cleanup_movs instructions =
   let is_redundant_mov = function
@@ -92,6 +93,7 @@ module Allocator (R : REG_TYPE) = struct
     String.replace_chars (function '.' -> "_" | c -> String.of_char c) s
 
   let k = Set.cardinal all_hardregs
+  let get_node_by_id graph node_id = Map.find node_id graph
 
   let add_edge g nd_id1 nd_id2 =
     let nd1 = Map.find nd_id1 g in
@@ -218,6 +220,19 @@ module Allocator (R : REG_TYPE) = struct
         dump_helper ".ncol" edge_printer post_processor ctx g
   end
 
+  let remove_edge g nd_id1 nd_id2 =
+    let nd1, nd2 = (get_node_by_id g nd_id1, get_node_by_id g nd_id2) in
+    nd1.neighbors <- Set.remove nd_id2 nd1.neighbors;
+    nd2.neighbors <- Set.remove nd_id1 nd2.neighbors
+
+  let degree graph nd_id =
+    let nd = get_node_by_id graph nd_id in
+    Set.cardinal nd.neighbors
+
+  let are_neighbors g nd_id1 nd_id2 =
+    let nd1 = Map.find nd_id1 g in
+    Set.mem nd_id2 nd1.neighbors
+
   module LivenessAnalysis = struct
     open AsmCfg
     module Iterative = Backward_dataflow.Dataflow (AsmCfg)
@@ -382,6 +397,98 @@ module Allocator (R : REG_TYPE) = struct
     in
     Map.map set_spill_cost graph
 
+  let george_test graph ~hardreg ~pseudo =
+    let pseudoreg_neighbors = (get_node_by_id graph pseudo).neighbors in
+    let neighbor_is_ok neighbor_id =
+      (* a neighbor of the pseudo won't interfere with coalescing
+       * if it has insignificant degree or it already interferes with hardreg *)
+      are_neighbors graph neighbor_id hardreg || degree graph neighbor_id < k
+    in
+    Set.for_all neighbor_is_ok pseudoreg_neighbors
+
+  let briggs_test graph x y =
+    let x_nd = get_node_by_id graph x in
+    let y_nd = get_node_by_id graph y in
+    let neighbors = Set.union x_nd.neighbors y_nd.neighbors in
+    let has_significant_degree neighbor_id =
+      let deg = degree graph neighbor_id in
+      let adjusted_deg =
+        if
+          are_neighbors graph x neighbor_id && are_neighbors graph y neighbor_id
+        then deg - 1
+        else deg
+      in
+      adjusted_deg >= k
+    in
+    let count_significant neighbor cnt =
+      if has_significant_degree neighbor then cnt + 1 else cnt
+    in
+    let significant_neighbor_count = Set.fold count_significant neighbors 0 in
+    significant_neighbor_count < k
+
+  let conservative_coalescable graph src dst =
+    if briggs_test graph src dst then true
+    else
+      match (src, dst) with
+      | Reg _, _ -> george_test graph ~hardreg:src ~pseudo:dst
+      | _, Reg _ -> george_test graph ~hardreg:dst ~pseudo:src
+      | _ -> false
+
+  (* debug output *)
+  let print_coalesce_msg ctx src dst =
+    if (R.debug_settings ()).debug_msg && Debug.is_dump_target ctx then
+      Printf.printf "Coalescing %s into %s\n" (show_node_id src)
+        (show_node_id dst)
+
+  let update_graph ctx g ~to_merge ~to_keep =
+    let update_neighbor neighbor_id =
+      add_edge g neighbor_id to_keep;
+      remove_edge g neighbor_id to_merge
+    in
+    print_coalesce_msg ctx to_merge to_keep;
+    Set.iter update_neighbor (get_node_by_id g to_merge).neighbors;
+    Map.remove to_merge g
+
+  let coalesce ctx graph instructions =
+    if (R.debug_settings ()).debug_msg && Debug.is_dump_target ctx then
+      Printf.printf "Coalescing round\n=============\n";
+    let process_instr (g, reg_map) = function
+      | Mov (_, src, dst) ->
+          let src' = Disjoint_sets.find src reg_map in
+          let dst' = Disjoint_sets.find dst reg_map in
+          if
+            Map.mem src' g
+            && Map.mem dst' g
+            && src' <> dst'
+            && (not (are_neighbors g src' dst'))
+            && conservative_coalescable g src' dst'
+          then
+            match src' with
+            | Reg _ ->
+                ( update_graph ctx g ~to_merge:dst' ~to_keep:src',
+                  Disjoint_sets.union dst' src' reg_map )
+            | _ ->
+                ( update_graph ctx g ~to_merge:src' ~to_keep:dst',
+                  Disjoint_sets.union src' dst' reg_map )
+          else (g, reg_map)
+      | _ -> (g, reg_map)
+    in
+    let _updated_graph, new_instructions =
+      List.fold process_instr (graph, Disjoint_sets.init) instructions
+    in
+    new_instructions
+
+  let rewrite_coalesced instructions coalesced_regs =
+    let f r = Disjoint_sets.find r coalesced_regs in
+    let rewrite_instruction = function
+      | Mov (t, src, dst) ->
+          let new_src = f src in
+          let new_dst = f dst in
+          if new_src = new_dst then None else Some (Mov (t, new_src, new_dst))
+      | i -> Some (replace_ops f i)
+    in
+    List.filter_map rewrite_instruction instructions
+
   let rec color_graph ctx graph =
     let remaining =
       graph
@@ -418,8 +525,7 @@ module Allocator (R : REG_TYPE) = struct
             let spill_info_printer fmt =
               Printf.ksprintf
                 (fun msg ->
-                  if
-                    (R.debug_settings ()).spill_info && Debug.is_dump_target ctx
+                  if (R.debug_settings ()).debug_msg && Debug.is_dump_target ctx
                   then Printf.printf "%s" msg)
                 fmt
             in
@@ -505,16 +611,28 @@ module Allocator (R : REG_TYPE) = struct
     cleanup_movs (List.map (replace_ops f) instructions)
 
   let allocate ctx aliased_pseudos instructions =
-    let graph : graph =
-      build_interference_graph ctx aliased_pseudos instructions
+    let rec coalesce_loop current_instructions =
+      let graph : graph =
+        build_interference_graph ctx aliased_pseudos current_instructions
+      in
+      DumpGraph.dump_graphviz ctx graph;
+      DumpGraph.dump_ncol ctx graph;
+      let coalesced_regs = coalesce ctx graph current_instructions in
+      if Disjoint_sets.is_empty coalesced_regs then (graph, current_instructions)
+      else
+        let new_instructions =
+          rewrite_coalesced current_instructions coalesced_regs
+        in
+        coalesce_loop new_instructions
+    in
+    let coalesced_graph, coalesced_instructions = coalesce_loop instructions in
+    let graph_with_spill_costs =
+      add_spill_costs coalesced_graph coalesced_instructions
     in
-    DumpGraph.dump_graphviz ctx graph;
-    DumpGraph.dump_ncol ctx graph;
-    let graph_with_spill_costs = add_spill_costs graph instructions in
     let colored_graph = color_graph ctx graph_with_spill_costs in
 
     let register_map = make_register_map ctx colored_graph in
-    replace_pseudoregs instructions register_map
+    replace_pseudoregs coalesced_instructions register_map
 end
 
 module GP = Allocator (struct

lib/dune

@@ -21,6 +21,7 @@
   copy_prop
   dead_store_elim
   debug
+  disjoint_sets
   emit
   extended_big_int
   optimize_utils

lib/settings.ml

@@ -31,7 +31,7 @@ type optimizations = {
 }
 
 type regalloc_debug_options = {
-  spill_info : bool;
+  debug_msg : bool;
   interference_ncol : bool;
   interference_graphviz : bool;
   liveness : bool;
@@ -55,7 +55,7 @@ type debug_options = {
 let debug =
   let regalloc_default =
     {
-      spill_info = false;
+      debug_msg = false;
       interference_ncol = false;
       interference_graphviz = false;
       liveness = false;

lib/settings.mli

@@ -9,7 +9,7 @@ type optimizations = {
 }
 
 type regalloc_debug_options = {
-  spill_info : bool;
+  debug_msg : bool;
   interference_ncol : bool;
   interference_graphviz : bool;
   liveness : bool;

lib/util/disjoint_sets.ml

@@ -0,0 +1,14 @@
+open Batteries
+
+type 'a t = ('a, 'a) Map.t
+
+let init = Map.empty
+let union x y disj_sets = Map.add x y disj_sets
+
+let rec find x disj_sets =
+  if Map.mem x disj_sets then
+    let mapped_to = Map.find x disj_sets in
+    find mapped_to disj_sets
+  else x
+
+let is_empty = Map.is_empty

lib/util/disjoint_sets.mli

@@ -0,0 +1,6 @@
+type 'a t
+
+val init : 'a t
+val union : 'a -> 'a -> 'a t -> 'a t
+val find : 'a -> 'a t -> 'a
+val is_empty : 'a t -> bool