Effects: partial CPS transform

CHANGES.md

@@ -2,7 +2,7 @@
 ## Features/Changes
 * Misc: bump min ocaml version to 4.08
 * Misc: remove some old runtime files to support some external libs
-* Effects: improved CPS transform, resulting in lower compilation time and smaller generated code
+* Effects: partial CPS transformation, resulting in much better performances, lower compilation time and smaller generated code
 * Compiler: separate compilation can now drops unused units when linking (similar to ocamlc). (#1378)
   Feature is disabled by default while dune rules are being fixed. Enable with --enable=auto-link.
 * Compiler: specialize string to js-string conversion for all valid utf8 strings (previously just ascii)

benchmarks/Makefile

@@ -158,7 +158,7 @@ time-effects.svg: __run_effects
 	 -omit minesweeper \
 	 -omit planet \
 	 -omit ocamlc \
-	 -max 5 -svg 7 400 150 -edgecaption -ylabel "Execution time" \
+	 -min 0.5 -max 1.5 -svg 7 400 150 -edgecaption -ylabel "Execution time" \
 	 > $@
 
 size-effects.svg: __run_effects
@@ -175,7 +175,7 @@ size-effects.svg: __run_effects
 	 -append planet \
 	 -append js_of_ocaml \
 	 -append ocamlc \
-	 -max 2 -svg 7 650 150 -edgecaption -ylabel Size \
+	 -min 0.8 -max 1.25 -svg 7 650 150 -edgecaption -ylabel Size \
 	 > $@
 
 size-gzipped-effects.svg: __run_effects

compiler/lib/code.ml

@@ -107,6 +107,8 @@ module Var : sig
     val set : 'a t -> key -> 'a -> unit
 
     val make : size -> 'a -> 'a t
+
+    val iter : (key -> 'a -> unit) -> 'a t -> unit
   end
 
   module ISet : sig
@@ -213,6 +215,11 @@ end = struct
     let set t x v = t.(x) <- v
 
     let make () v = Array.make (count ()) v
+
+    let iter f t =
+      for i = 0 to Array.length t - 1 do
+        f i t.(i)
+      done
   end
 
   module ISet = struct
@@ -633,6 +640,25 @@ let rec preorder_traverse' { fold } f pc visited blocks acc =
 let preorder_traverse fold f pc blocks acc =
   snd (preorder_traverse' fold f pc Addr.Set.empty blocks acc)
 
+let fold_closures_innermost_first { start; blocks; _ } f accu =
+  let rec visit blocks pc f accu =
+    traverse
+      { fold = fold_children }
+      (fun pc accu ->
+        let block = Addr.Map.find pc blocks in
+        List.fold_left block.body ~init:accu ~f:(fun accu i ->
+            match i with
+            | Let (x, Closure (params, cont)) ->
+                let accu = visit blocks (fst cont) f accu in
+                f (Some x) params cont accu
+            | _ -> accu))
+      pc
+      blocks
+      accu
+  in
+  let accu = visit blocks start f accu in
+  f None [] (start, []) accu
+
 let eq p1 p2 =
   p1.start = p2.start
   && Addr.Map.cardinal p1.blocks = Addr.Map.cardinal p2.blocks

compiler/lib/code.mli

@@ -100,6 +100,8 @@ module Var : sig
     val set : 'a t -> key -> 'a -> unit
 
     val make : size -> 'a -> 'a t
+
+    val iter : (key -> 'a -> unit) -> 'a t -> unit
   end
 
   module ISet : sig
@@ -231,6 +233,9 @@ type fold_blocs_poly = { fold : 'a. 'a fold_blocs } [@@unboxed]
 val fold_closures :
   program -> (Var.t option -> Var.t list -> cont -> 'd -> 'd) -> 'd -> 'd
 
+val fold_closures_innermost_first :
+  program -> (Var.t option -> Var.t list -> cont -> 'd -> 'd) -> 'd -> 'd
+
 val fold_children : 'c fold_blocs
 
 val traverse :

compiler/lib/dgraph.ml

@@ -58,22 +58,22 @@ struct
 
     let m = ref 0
 
-    type stack =
-      { stack : N.t Stack.t
+    type queue =
+      { queue : N.t Queue.t
       ; mutable set : NSet.t
       }
 
-    let is_empty st = Stack.is_empty st.stack
+    let is_empty st = Queue.is_empty st.queue
 
     let pop st =
-      let x = Stack.pop st.stack in
+      let x = Queue.pop st.queue in
       st.set <- NSet.remove x st.set;
       x
 
     let push x st =
       if not (NSet.mem x st.set)
       then (
-        Stack.push x st.stack;
+        Queue.push x st.queue;
         st.set <- NSet.add x st.set)
 
     let rec iterate g f v w =
@@ -91,24 +91,26 @@ struct
           iterate g f v w)
         else iterate g f v w
 
-    let rec traverse g visited stack x =
+    let rec traverse g visited lst x =
       if not (NSet.mem x visited)
       then (
         let visited = NSet.add x visited in
         let visited =
-          g.fold_children (fun y visited -> traverse g visited stack y) x visited
+          g.fold_children (fun y visited -> traverse g visited lst y) x visited
         in
-        Stack.push x stack;
+        lst := x :: !lst;
         visited)
       else visited
 
     let traverse_all g =
-      let stack = Stack.create () in
+      let lst = ref [] in
       let visited =
-        NSet.fold (fun x visited -> traverse g visited stack x) g.domain NSet.empty
+        NSet.fold (fun x visited -> traverse g visited lst x) g.domain NSet.empty
       in
       assert (NSet.equal g.domain visited);
-      stack
+      let queue = Queue.create () in
+      List.iter ~f:(fun x -> Queue.push x queue) !lst;
+      queue
 
     let f g f =
       n := 0;
@@ -128,7 +130,7 @@ let t1 = Timer.make () in
 let t1 = Timer.get t1 in
 let t2 = Timer.make () in
 *)
-      let w = { set = g.domain; stack = traverse_all g } in
+      let w = { set = g.domain; queue = traverse_all g } in
       (*
 let t2 = Timer.get t2 in
 let t3 = Timer.make () in
@@ -206,54 +208,55 @@ struct
 
     let m = ref 0
 
-    type stack =
-      { stack : N.t Stack.t
+    type queue =
+      { queue : N.t Queue.t
       ; set : NSet.t
       }
 
-    let is_empty st = Stack.is_empty st.stack
+    let is_empty st = Queue.is_empty st.queue
 
     let pop st =
-      let x = Stack.pop st.stack in
+      let x = Queue.pop st.queue in
       NSet.add st.set x;
       x
 
     let push x st =
       if NSet.mem st.set x
       then (
-        Stack.push x st.stack;
+        Queue.push x st.queue;
         NSet.remove st.set x)
 
-    let rec iterate g f v w =
+    let rec iterate g ~update f v w =
       if is_empty w
       then v
       else
         let x = pop w in
         let a = NTbl.get v x in
         incr m;
-        let b = f v x in
-        NTbl.set v x b;
+        let b = f ~update v x in
         if not (D.equal a b)
         then (
-          g.iter_children (fun y -> push y w) x;
-          iterate g f v w)
-        else iterate g f v w
+          NTbl.set v x b;
+          g.iter_children (fun y -> push y w) x);
+        iterate g ~update f v w
 
-    let rec traverse g to_visit stack x =
+    let rec traverse g to_visit lst x =
       if NSet.mem to_visit x
       then (
         NSet.remove to_visit x;
         incr n;
-        g.iter_children (fun y -> traverse g to_visit stack y) x;
-        Stack.push x stack)
+        g.iter_children (fun y -> traverse g to_visit lst y) x;
+        lst := x :: !lst)
 
     let traverse_all g =
-      let stack = Stack.create () in
+      let lst = ref [] in
       let to_visit = NSet.copy g.domain in
-      NSet.iter (fun x -> traverse g to_visit stack x) g.domain;
-      { stack; set = to_visit }
+      NSet.iter (fun x -> traverse g to_visit lst x) g.domain;
+      let queue = Queue.create () in
+      List.iter ~f:(fun x -> Queue.push x queue) !lst;
+      { queue; set = to_visit }
 
-    let f size g f =
+    let f' size g f =
       n := 0;
       m := 0;
       (*
@@ -269,12 +272,17 @@ let t2 = Timer.make () in
 let t2 = Timer.get t2 in
 let t3 = Timer.make () in
 *)
-      let res = iterate g f v w in
+      let update ~children x =
+        if children then g.iter_children (fun y -> push y w) x else push x w
+      in
+      let res = iterate g ~update f v w in
       (*
 let t3 = Timer.get t3 in
       Format.eprintf "YYY %.2f %.2f %.2f@." t1 t2 t3;
       Format.eprintf "YYY %d %d (%f)@." !m !n (float !m /. float !n);
 *)
       res
+
+    let f size g f = f' size g (fun ~update:_ v x -> f v x)
   end
 end

compiler/lib/dgraph.mli

@@ -94,5 +94,11 @@ end)
 
   module Solver (D : DOMAIN) : sig
     val f : NTbl.size -> t -> (D.t NTbl.t -> N.t -> D.t) -> D.t NTbl.t
+
+    val f' :
+         NTbl.size
+      -> t
+      -> (update:(children:bool -> N.t -> unit) -> D.t NTbl.t -> N.t -> D.t)
+      -> D.t NTbl.t
   end
 end

compiler/lib/driver.ml

@@ -28,11 +28,8 @@ let should_export = function
   | `Named _ | `Anonymous -> true
 
 let tailcall p =
-  if Config.Flag.effects ()
-  then p
-  else (
-    if debug () then Format.eprintf "Tail-call optimization...@.";
-    Tailcall.f p)
+  if debug () then Format.eprintf "Tail-call optimization...@.";
+  Tailcall.f p
 
 let deadcode' p =
   if debug () then Format.eprintf "Dead-code...@.";
@@ -83,19 +80,21 @@ let phi p =
   if debug () then Format.eprintf "Variable passing simplification...@.";
   Phisimpl.f p
 
+let ( +> ) f g x = g (f x)
+
+let map_fst f (x, y) = f x, y
+
 let effects p =
   if Config.Flag.effects ()
   then (
     if debug () then Format.eprintf "Effects...@.";
-    Deadcode.f p |> Effects.f |> Lambda_lifting.f)
-  else p
+    p |> Deadcode.f +> Effects.f +> map_fst Lambda_lifting.f)
+  else p, (Code.Var.Set.empty : Effects.cps_calls)
 
 let print p =
   if debug () then Code.Print.program (fun _ _ -> "") p;
   p
 
-let ( +> ) f g x = g (f x)
-
 let rec loop max name round i (p : 'a) : 'a =
   let p' = round p in
   if i >= max || Code.eq p' p
@@ -154,10 +153,22 @@ let round2 = flow +> specialize' +> eval +> deadcode +> o1
 
 let o3 = loop 10 "tailcall+inline" round1 1 +> loop 10 "flow" round2 1 +> print
 
-let generate d ~exported_runtime ~wrap_with_fun ~warn_on_unhandled_effect (p, live_vars) =
+let generate
+    d
+    ~exported_runtime
+    ~wrap_with_fun
+    ~warn_on_unhandled_effect
+    ((p, live_vars), cps_calls) =
   if times () then Format.eprintf "Start Generation...@.";
   let should_export = should_export wrap_with_fun in
-  Generate.f p ~exported_runtime ~live_vars ~should_export ~warn_on_unhandled_effect d
+  Generate.f
+    p
+    ~exported_runtime
+    ~live_vars
+    ~cps_calls
+    ~should_export
+    ~warn_on_unhandled_effect
+    d
 
 let header formatter ~custom_header =
   match custom_header with
@@ -553,7 +564,9 @@ let full
     d
     p =
   let exported_runtime = not standalone in
-  let opt = specialize_js_once +> profile +> effects +> Generate_closure.f +> deadcode' in
+  let opt =
+    specialize_js_once +> profile +> effects +> map_fst (Generate_closure.f +> deadcode')
+  in
   let emit =
     generate d ~exported_runtime ~wrap_with_fun ~warn_on_unhandled_effect:standalone
     +> link ~standalone ~linkall