Register allocation: try aggressive coalescing first

backend/IRC.ml

@@ -669,7 +669,7 @@ let combine g u v =
 
 (* Attempt coalescing *)
 
-let coalesce g =
+let coalesce forced_coalescing g =
   let m = DLinkMove.pick g.worklistMoves in
   let x = getAlias m.src and y = getAlias m.dst in
   let (u, v) = if y.nstate = Colored then (y, x) else (x, y) in
@@ -681,7 +681,7 @@ let coalesce g =
     DLinkMove.insert m g.constrainedMoves;
     addWorkList g u;
     addWorkList g v
-  end else if canCoalesce g u v then begin
+  end else if forced_coalescing || canCoalesce g u v then begin
     DLinkMove.insert m g.coalescedMoves;
     combine g u v;
     addWorkList g u
@@ -758,16 +758,16 @@ let selectSpill g =
 
 (* Produce the order of nodes that we'll use for coloring *)
 
-let rec nodeOrder g stack =
+let rec nodeOrder forced g stack =
   (*i checkInvariants g; *)
   if DLinkNode.notempty g.simplifyWorklist then
-    (let n = simplify g in nodeOrder g (n :: stack))
+    (let n = simplify g in nodeOrder forced g (n :: stack))
   else if DLinkMove.notempty g.worklistMoves then
-    (coalesce g; nodeOrder g stack)
+    (coalesce forced g; nodeOrder forced g stack)
   else if DLinkNode.notempty g.freezeWorklist then
-    (freeze g; nodeOrder g stack)
+    (freeze g; nodeOrder forced g stack)
   else if DLinkNode.notempty g.spillWorklist then
-    (let n = selectSpill g in nodeOrder g (n :: stack))
+    (let n = selectSpill g in nodeOrder forced g (n :: stack))
   else
     stack
 
@@ -921,7 +921,8 @@ let add_interf g v1 v2 =
 let add_pref g v1 v2 =
   let n1 = nodeOfVar g v1 in let n2 = nodeOfVar g v2 in addMovePref g n1 n2
 
-let coloring g =
+let coloring ~coalescing g =
+  let forced = match coalescing with `Aggressive -> true | `Prudent -> false in
   initialNodePartition g;
-  List.iter (assign_color g) (nodeOrder g []);
+  List.iter (assign_color g) (nodeOrder forced g []);
   location_of_var g  (* total function var -> location *)

backend/IRC.mli

@@ -35,8 +35,12 @@ val add_interf: graph -> var -> var -> unit
 (* Add a preference between two variables. *)
 val add_pref: graph -> var -> var -> unit
 
-(* Color the graph.  Return an assignment of locations to variables. *)
-val coloring: graph -> (var -> loc)
+(* Color the graph.  Return an assignment of locations to variables.
+   [~coalescing] sets the coalescing mode: 
+   - [`Prudent] to perform only coalescings that cannot cause spilling;
+   - [`Aggressive] to force coalescing whenever semantically sound. *)
+
+val coloring: coalescing:[`Prudent|`Aggressive] -> graph -> (var -> loc)
 
 (* Auxiliaries to deal with register classes *)
 val class_of_type: AST.typ -> int

backend/Regalloc.ml

@@ -733,7 +733,7 @@ let rec add_interfs_block g blk live =
       add_interfs_instr g instr live';
       live_before instr live'
 
-let find_coloring f liveness =
+let find_coloring ~coalescing f liveness =
   (*type_function f;  (* for debugging *)*)
   let g = IRC.init (spill_costs f) in
   PTree.fold
@@ -742,7 +742,7 @@ let find_coloring f liveness =
   add_interfs_destroyed g
     (transfer_live f f.fn_entrypoint (PMap.get f.fn_entrypoint liveness))
     destroyed_at_function_entry;
-  IRC.coloring g
+  IRC.coloring ~coalescing g
 
 
 (*********** Determination of variables that need spill code insertion *****)
@@ -1153,19 +1153,28 @@ let transl_function fn alloc =
 exception Timeout
 
 let rec first_round f liveness =
-  let alloc = find_coloring f liveness in
+  let alloc = find_coloring ~coalescing:`Aggressive f liveness in
   if !option_dalloctrace then begin
     fprintf !pp "-------------- After initial register allocation\n\n";
     PrintXTL.print_function !pp ~alloc: alloc ~live: liveness f
   end;
   let ts = tospill_function f alloc in
+  if VSet.is_empty ts then success f alloc else second_round f liveness
+
+and second_round f liveness =
+  let alloc = find_coloring ~coalescing:`Prudent f liveness in
+  if !option_dalloctrace then begin
+    fprintf !pp "-------------- After register allocation without aggressive coalescing\n\n";
+    PrintXTL.print_function !pp ~alloc: alloc ~live: liveness f
+  end;
+  let ts = tospill_function f alloc in
   if VSet.is_empty ts then success f alloc else more_rounds f ts 1
 
 and more_rounds f ts count =
   if count >= 40 then raise Timeout;
   let f' = spill_function f ts count in
   let liveness = liveness_analysis f' in
-  let alloc = find_coloring f' liveness in
+  let alloc = find_coloring ~coalescing:`Prudent f' liveness in
   if !option_dalloctrace then begin
     fprintf !pp "-------------- After register allocation (round %d)\n\n" count;
     PrintXTL.print_function !pp ~alloc: alloc ~live: liveness f'