Reference unboxing

compiler/lib/deadcode.ml

@@ -183,6 +183,26 @@ let annot st pc xi =
 
 (****)
 
+let remove_unused_blocks p =
+  let visited = BitSet.create' p.free_pc in
+  let rec mark_used pc =
+    if not (BitSet.mem visited pc)
+    then (
+      BitSet.set visited pc;
+      let block = Addr.Map.find pc p.blocks in
+      List.iter
+        ~f:(fun i ->
+          match i with
+          | Let (_, Closure (_, (pc', _), _)) -> mark_used pc'
+          | _ -> ())
+        block.body;
+      Code.fold_children p.blocks pc (fun pc' () -> mark_used pc') ())
+  in
+  mark_used p.start;
+  { p with blocks = Addr.Map.filter (fun pc _ -> BitSet.mem visited pc) p.blocks }
+
+(****)
+
 let rec add_arg_dep defs params args =
   match params, args with
   | x :: params, y :: args ->
@@ -250,7 +270,7 @@ let remove_empty_blocks ~live_vars (p : Code.program) : Code.program =
         })
       p.blocks
   in
-  { p with blocks }
+  remove_unused_blocks { p with blocks }
 
 let f ({ blocks; _ } as p : Code.program) =
   let t = Timer.make () in

compiler/lib/deadcode.mli

@@ -24,3 +24,5 @@ type variable_uses =
 val f : Code.program -> Code.program * variable_uses
 
 val remove_empty_blocks : live_vars:variable_uses -> Code.program -> Code.program
+
+val remove_unused_blocks : Code.program -> Code.program

compiler/lib/driver.ml

@@ -178,6 +178,7 @@ let o1 : 'a -> 'a =
   +> flow
   +> specialize
   +> eval
+  +> Ref_unboxing.f
   +> inline
   +> deadcode
   +> phi

compiler/lib/eval.ml

@@ -760,11 +760,12 @@ let f info p =
   let t = Timer.make () in
   let blocks = eval update_count ~target:(Config.target ()) info p.blocks in
   let blocks = drop_exception_handler drop_count blocks in
+  let p = Deadcode.remove_unused_blocks { p with blocks } in
   if times () then Format.eprintf "  eval: %a@." Timer.print t;
   if stats ()
   then
     Format.eprintf
       "Stats - eval: %d optimizations, %d dropped exception handlers@."
       !update_count
       !drop_count;
-  { p with blocks }
+  p

compiler/lib/phisimpl.ml

@@ -67,6 +67,10 @@ let program_deps { blocks; _ } =
     (fun _pc block ->
       List.iter block.body ~f:(fun i ->
           match i with
+          | Let (x, Prim (Extern "%identity", [ Pv y ])) ->
+              add_var vars x;
+              add_dep deps x y;
+              add_def vars defs x y
           | Let (x, e) ->
               add_var vars x;
               expr_deps blocks vars deps defs x e

compiler/lib/ref_unboxing.ml

@@ -0,0 +1,166 @@
+open! Stdlib
+open Code
+
+(*
+ocamlc does not perform reference unboxing when emitting debugging
+information. Inlining can enable additional reference unboxing.
+
+TODO: handle assignment in handler
+*)
+
+let debug = Debug.find "unbox-refs"
+
+let times = Debug.find "times"
+
+let stats = Debug.find "stats"
+
+let rewrite refs block m =
+  let m, l =
+    List.fold_left
+      ~f:(fun (m, rem) i ->
+        match i with
+        | Let (x, Block (0, [| y |], (NotArray | Unknown), Maybe_mutable))
+          when Var.Set.mem x refs -> Var.Map.add x y m, rem
+        | Let (y, Field (x, 0, Non_float)) when Var.Map.mem x m ->
+            m, Let (y, Prim (Extern "%identity", [ Pv (Var.Map.find x m) ])) :: rem
+        | Offset_ref (x, n) when Var.Map.mem x m ->
+            let y = Var.fresh () in
+            ( Var.Map.add x y m
+            , Let
+                ( y
+                , Prim
+                    ( Extern "%int_add"
+                    , [ Pv (Var.Map.find x m); Pc (Int (Targetint.of_int_exn n)) ] ) )
+              :: rem )
+        | Set_field (x, _, Non_float, y) when Var.Map.mem x m -> Var.Map.add x y m, rem
+        | _ -> m, i :: rem)
+      block.body
+      ~init:(m, [])
+  in
+  m, List.rev l
+
+let rewrite_cont relevant_vars vars (pc', args) =
+  let refs, _ = Hashtbl.find relevant_vars pc' in
+  let vars = Var.Map.filter (fun x _ -> Var.Set.mem x refs) vars in
+  pc', List.map ~f:snd (Var.Map.bindings vars) @ args
+
+let rewrite_function p variables pc =
+  let relevant_vars = Hashtbl.create 16 in
+  let g = Structure.(dominator_tree (build_graph p.blocks pc)) in
+  let rec traverse_tree g pc vars =
+    let block = Addr.Map.find pc p.blocks in
+    let vars' =
+      List.fold_left
+        ~f:(fun s i ->
+          match i with
+          | Let (x, Block (0, [| _ |], (NotArray | Unknown), Maybe_mutable))
+            when Var.Hashtbl.mem variables x -> Var.Set.add x s
+          | _ -> s)
+        ~init:vars
+        block.body
+    in
+    Hashtbl.add relevant_vars pc (vars, vars');
+    Addr.Set.iter (fun pc' -> traverse_tree g pc' vars') (Structure.get_edges g pc)
+  in
+  traverse_tree g pc Var.Set.empty;
+  let rec traverse_tree' g pc blocks =
+    let block = Addr.Map.find pc p.blocks in
+    let vars, refs = Hashtbl.find relevant_vars pc in
+    let vars =
+      Var.Set.fold (fun x m -> Var.Map.add x (Var.fork x) m) vars Var.Map.empty
+    in
+    let params = List.map ~f:snd (Var.Map.bindings vars) @ block.params in
+    let vars, body = rewrite refs block vars in
+    let branch =
+      match block.branch with
+      | Return _ | Raise _ | Stop -> block.branch
+      | Branch cont -> Branch (rewrite_cont relevant_vars vars cont)
+      | Cond (x, cont, cont') ->
+          Cond
+            ( x
+            , rewrite_cont relevant_vars vars cont
+            , rewrite_cont relevant_vars vars cont' )
+      | Switch (x, a) ->
+          Switch (x, Array.map ~f:(fun cont -> rewrite_cont relevant_vars vars cont) a)
+      | Pushtrap (cont, x, cont') ->
+          Pushtrap
+            ( rewrite_cont relevant_vars vars cont
+            , x
+            , rewrite_cont relevant_vars vars cont' )
+      | Poptrap cont -> Poptrap (rewrite_cont relevant_vars vars cont)
+    in
+    let blocks = Addr.Map.add pc { params; body; branch } blocks in
+    Addr.Set.fold
+      (fun pc' blocks -> traverse_tree' g pc' blocks)
+      (Structure.get_edges g pc)
+      blocks
+  in
+  let blocks = traverse_tree' g pc p.blocks in
+  { p with blocks }
+
+let f p =
+  let t = Timer.make () in
+  let candidates = Var.Hashtbl.create 128 in
+  let updated = Var.Hashtbl.create 128 in
+  let visited = BitSet.create' p.free_pc in
+  let discard x = Var.Hashtbl.remove candidates x in
+  let check_field_access depth x =
+    match Var.Hashtbl.find candidates x with
+    | exception Not_found -> false
+    | depth' ->
+        if depth' = depth
+        then true
+        else (
+          Var.Hashtbl.remove candidates x;
+          false)
+  in
+  let rec traverse depth start_pc pc =
+    if not (BitSet.mem visited pc)
+    then (
+      BitSet.set visited pc;
+      let block = Addr.Map.find pc p.blocks in
+      List.iter
+        ~f:(fun i ->
+          match i with
+          | Let (x, Block (0, [| _ |], (NotArray | Unknown), Maybe_mutable)) ->
+              Freevars.iter_instr_free_vars discard i;
+              Var.Hashtbl.replace candidates x depth
+          | Let (_, Closure (_, (pc', _), _)) -> traverse (depth + 1) pc' pc'
+          | Let (_, Field (x, 0, Non_float)) -> ignore (check_field_access depth x)
+          | Offset_ref (x, _) ->
+              if check_field_access depth x then Var.Hashtbl.replace updated x start_pc
+          | Set_field (x, _, Non_float, y) ->
+              discard y;
+              if check_field_access depth x then Var.Hashtbl.replace updated x start_pc
+          | _ -> Freevars.iter_instr_free_vars discard i)
+        block.body;
+      Freevars.iter_last_free_var discard block.branch;
+      match block.branch with
+      | Pushtrap ((pc', _), _, (pc'', _)) ->
+          traverse (depth + 1) start_pc pc';
+          traverse depth start_pc pc''
+      | Poptrap (pc', _) -> traverse (depth - 1) start_pc pc'
+      | _ -> Code.fold_children p.blocks pc (fun pc' () -> traverse depth start_pc pc') ())
+  in
+  traverse 0 p.start p.start;
+  if debug ()
+  then
+    Print.program
+      Format.err_formatter
+      (fun _ i ->
+        match i with
+        | Instr (Let (x, _))
+          when Var.Hashtbl.mem candidates x && Var.Hashtbl.mem updated x -> "REF"
+        | _ -> "")
+      p;
+  Var.Hashtbl.filter_map_inplace
+    (fun x _depth -> try Some (Var.Hashtbl.find updated x) with Not_found -> None)
+    candidates;
+  let functions =
+    Var.Hashtbl.fold (fun _ pc s -> Addr.Set.add pc s) candidates Addr.Set.empty
+  in
+  let p = Addr.Set.fold (fun pc p -> rewrite_function p candidates pc) functions p in
+  if times () then Format.eprintf "  reference unboxing: %a@." Timer.print t;
+  if stats ()
+  then Format.eprintf "Stats - reference unboxing: %d@." (Var.Hashtbl.length candidates);
+  p