Compiler: make tailcall optim more robust

CHANGES.md

@@ -26,6 +26,7 @@
 * Compiler: static evaluation of more primitives (#1912)
 * Compiler: faster compilation by stopping sooner when optimizations become unproductive (#1939)
 * Compiler: improve debug/sourcemap location of closures (#1947)
+* Compiler: improve tailcall optimization (#1943)
 * Runtime: use Dataview to convert between floats and bit representation
 
 ## Bug fixes

compiler/lib/driver.ml

@@ -70,9 +70,13 @@ let specialize_js (p, info) =
   if debug () then Format.eprintf "Specialize js...@.";
   Specialize_js.f info p
 
-let specialize_js_once p =
+let specialize_js_once_before p =
   if debug () then Format.eprintf "Specialize js once...@.";
-  Specialize_js.f_once p
+  Specialize_js.f_once_before p
+
+let specialize_js_once_after p =
+  if debug () then Format.eprintf "Specialize js once...@.";
+  Specialize_js.f_once_after p
 
 let specialize' (p, info) =
   let p = specialize_1 (p, info) in
@@ -87,10 +91,6 @@ let flow p =
   if debug () then Format.eprintf "Data flow...@.";
   Flow.f p
 
-let flow_simple p =
-  if debug () then Format.eprintf "Data flow...@.";
-  Flow.f ~skip_param:true p
-
 let phi p =
   if debug () then Format.eprintf "Variable passing simplification...@.";
   Phisimpl.f p
@@ -161,7 +161,7 @@ let identity x = x
 let o1 : 'a -> 'a =
   print
   +> tailcall
-  +> flow_simple (* flow simple to keep information for future tailcall opt *)
+  +> flow
   +> specialize'
   +> eval
   +> inline (* inlining may reveal new tailcall opt *)
@@ -190,19 +190,7 @@ let o2 : 'a -> 'a = loop 10 "o1" o1 1 +> print
 
 (* o3 *)
 
-let round1 : 'a -> 'a =
-  print
-  +> tailcall
-  +> inline (* inlining may reveal new tailcall opt *)
-  +> deadcode (* deadcode required before flow simple -> provided by constant *)
-  +> flow_simple (* flow simple to keep information for future tailcall opt *)
-  +> specialize'
-  +> eval
-  +> identity
-
-let round2 = flow +> specialize' +> eval +> deadcode +> o1
-
-let o3 = loop 10 "tailcall+inline" round1 1 +> loop 10 "flow" round2 1 +> print
+let o3 = loop 10 "o1" o1 1 +> print
 
 let generate
     ~exported_runtime
@@ -694,11 +682,12 @@ let optimize ~profile p =
   in
   let opt =
     Specialize.switches
-    +> specialize_js_once
+    +> specialize_js_once_before
     +> (match profile with
        | O1 -> o1
        | O2 -> o2
        | O3 -> o3)
+    +> specialize_js_once_after
     +> exact_calls ~deadcode_sentinal profile
     +> effects ~deadcode_sentinal
     +> map_fst

compiler/lib/flow.ml

@@ -291,9 +291,9 @@ let program_escape defs known_origins { blocks; _ } =
 
 (****)
 
-let propagate2 ?(skip_param = false) defs known_origins possibly_mutable st x =
+let propagate2 defs known_origins possibly_mutable st x =
   match defs.(Var.idx x) with
-  | Param -> skip_param
+  | Param -> false
   | Phi s -> Var.Set.exists (fun y -> Var.Tbl.get st y) s
   | Expr e -> (
       match e with
@@ -320,11 +320,11 @@ end
 
 module Solver2 = G.Solver (Domain2)
 
-let solver2 ?skip_param vars deps defs known_origins possibly_mutable =
+let solver2 vars deps defs known_origins possibly_mutable =
   let g =
     { G.domain = vars; G.iter_children = (fun f x -> Var.Set.iter f deps.(Var.idx x)) }
   in
-  Solver2.f () g (propagate2 ?skip_param defs known_origins possibly_mutable)
+  Solver2.f () g (propagate2 defs known_origins possibly_mutable)
 
 let get_approx
     { Info.info_defs = _; info_known_origins; info_maybe_unknown; _ }
@@ -501,7 +501,7 @@ let print_stats s =
   done;
   Format.eprintf "Stats - flow updates: %d@." !count
 
-let f ?skip_param p =
+let f p =
   Code.invariant p;
   let t = Timer.make () in
   let t1 = Timer.make () in
@@ -514,7 +514,7 @@ let f ?skip_param p =
   let possibly_mutable = program_escape defs known_origins p in
   if times () then Format.eprintf "    flow analysis 3: %a@." Timer.print t3;
   let t4 = Timer.make () in
-  let maybe_unknown = solver2 ?skip_param vars deps defs known_origins possibly_mutable in
+  let maybe_unknown = solver2 vars deps defs known_origins possibly_mutable in
   if times () then Format.eprintf "    flow analysis 4: %a@." Timer.print t4;
   if debug ()
   then

compiler/lib/flow.mli

@@ -65,4 +65,4 @@ val the_block_contents_of : Info.t -> Code.prim_arg -> Code.Var.t array option
 
 val the_int : Info.t -> Code.prim_arg -> Targetint.t option
 
-val f : ?skip_param:bool -> Code.program -> Code.program * Info.t
+val f : Code.program -> Code.program * Info.t

compiler/lib/inline.ml

@@ -166,13 +166,7 @@ let rewrite_closure blocks cont_pc clos_pc =
 
 (****)
 
-let rec args_equal xs ys =
-  match xs, ys with
-  | [], [] -> true
-  | x :: xs, Pv y :: ys -> Code.Var.compare x y = 0 && args_equal xs ys
-  | _ -> false
-
-let inline ~first_class_primitives live_vars closures pc (outer, p) =
+let inline live_vars closures pc (outer, p) =
   let block = Addr.Map.find pc p.blocks in
   let body, (outer, branch, p) =
     List.fold_right
@@ -300,24 +294,6 @@ let inline ~first_class_primitives live_vars closures pc (outer, p) =
                     let outer = { outer with size = outer.size + f_size } in
                     [], (outer, Branch (fresh_addr, args), { p with blocks; free_pc })
               | _ -> i :: rem, state)
-        | Let (x, Closure (l, (pc, []), _)) when first_class_primitives -> (
-            let block = Addr.Map.find pc p.blocks in
-            match block with
-            | { body =
-                  ( [ Let (y, Prim (Extern prim, args)) ]
-                  | [ Event _; Let (y, Prim (Extern prim, args)) ]
-                  | [ Event _; Let (y, Prim (Extern prim, args)); Event _ ] )
-              ; branch = Return y'
-              ; params = []
-              } ->
-                let len = List.length l in
-                if
-                  Code.Var.compare y y' = 0
-                  && Primitive.has_arity prim len
-                  && args_equal l args
-                then Let (x, Special (Alias_prim prim)) :: rem, state
-                else i :: rem, state
-            | _ -> i :: rem, state)
         | _ -> i :: rem, state)
   in
   outer, { p with blocks = Addr.Map.add pc { block with body; branch } p.blocks }
@@ -327,12 +303,6 @@ let inline ~first_class_primitives live_vars closures pc (outer, p) =
 let times = Debug.find "times"
 
 let f p live_vars =
-  let first_class_primitives =
-    match Config.target (), Config.effects () with
-    | `JavaScript, `Disabled -> true
-    | `JavaScript, (`Cps | `Double_translation) | `Wasm, _ -> false
-    | `JavaScript, `Jspi -> assert false
-  in
   Code.invariant p;
   let t = Timer.make () in
   let closures = get_closures p in
@@ -343,7 +313,7 @@ let f p live_vars =
         let traverse outer =
           Code.traverse
             { fold = Code.fold_children }
-            (inline ~first_class_primitives live_vars closures)
+            (inline live_vars closures)
             pc
             p.blocks
             (outer, p)

compiler/lib/specialize_js.ml

@@ -354,7 +354,7 @@ let specialize_all_instrs ~target info p =
 
 let f info p = specialize_all_instrs ~target:(Config.target ()) info p
 
-let f_once p =
+let f_once_before p =
   let rec loop acc l =
     match l with
     | [] -> List.rev acc
@@ -381,3 +381,47 @@ let f_once p =
     Addr.Map.map (fun block -> { block with Code.body = loop [] block.body }) p.blocks
   in
   { p with blocks }
+
+let rec args_equal xs ys =
+  match xs, ys with
+  | [], [] -> true
+  | x :: xs, Pv y :: ys -> Code.Var.compare x y = 0 && args_equal xs ys
+  | _ -> false
+
+let f_once_after p =
+  let first_class_primitives =
+    match Config.target (), Config.effects () with
+    | `JavaScript, `Disabled -> true
+    | `JavaScript, (`Cps | `Double_translation) | `Wasm, _ -> false
+    | `JavaScript, `Jspi -> assert false
+  in
+  let f = function
+    | Let (x, Closure (l, (pc, []), _)) as i -> (
+        let block = Addr.Map.find pc p.blocks in
+        match block with
+        | { body =
+              ( [ Let (y, Prim (Extern prim, args)) ]
+              | [ Event _; Let (y, Prim (Extern prim, args)) ]
+              | [ Event _; Let (y, Prim (Extern prim, args)); Event _ ] )
+          ; branch = Return y'
+          ; params = []
+          } ->
+            let len = List.length l in
+            if
+              Code.Var.compare y y' = 0
+              && Primitive.has_arity prim len
+              && args_equal l args
+            then Let (x, Special (Alias_prim prim))
+            else i
+        | _ -> i)
+    | i -> i
+  in
+  if first_class_primitives
+  then
+    let blocks =
+      Addr.Map.map
+        (fun block -> { block with Code.body = List.map block.body ~f })
+        p.blocks
+    in
+    { p with blocks }
+  else p

compiler/lib/specialize_js.mli

@@ -20,4 +20,6 @@
 
 val f : Flow.Info.t -> Code.program -> Code.program
 
-val f_once : Code.program -> Code.program
+val f_once_before : Code.program -> Code.program
+
+val f_once_after : Code.program -> Code.program