Disable tail-call optimization for Math builtins

V8 currently generates suboptimal code for this (the builtins are no longer
inlined).

Author: vouillon

compiler/lib-wasm/code_generation.ml

@@ -59,6 +59,7 @@ type context =
   ; mutable globalized_variables : Var.Set.t
   ; value_type : W.value_type
   ; mutable unit_name : string option
+  ; mutable no_tail_call : unit Var.Hashtbl.t
   }
 
 let make_context ~value_type =
@@ -82,6 +83,7 @@ let make_context ~value_type =
   ; globalized_variables = Var.Set.empty
   ; value_type
   ; unit_name = None
+  ; no_tail_call = Var.Hashtbl.create 16
   }
 
 type var =
@@ -224,27 +226,30 @@ let get_global name =
     | Some { init; _ } -> init
     | _ -> None)
 
-let register_import ?(import_module = "env") ~name typ st =
-  ( (try
-       let x, typ' =
-         StringMap.find name (StringMap.find import_module st.context.imports)
-       in
-       (*ZZZ error message*)
-       assert (Poly.equal typ typ');
-       x
-     with Not_found ->
-       let x = Var.fresh_n name in
-       st.context.imports <-
-         StringMap.update
-           import_module
-           (fun m ->
-             Some
-               (match m with
-               | None -> StringMap.singleton name (x, typ)
-               | Some m -> StringMap.add name (x, typ) m))
-           st.context.imports;
-       x)
-  , st )
+let register_import ?(allow_tail_call = true) ?(import_module = "env") ~name typ st =
+  let x =
+    try
+      let x, typ' =
+        StringMap.find name (StringMap.find import_module st.context.imports)
+      in
+      (*ZZZ error message*)
+      assert (Poly.equal typ typ');
+      x
+    with Not_found ->
+      let x = Var.fresh_n name in
+      st.context.imports <-
+        StringMap.update
+          import_module
+          (fun m ->
+            Some
+              (match m with
+              | None -> StringMap.singleton name (x, typ)
+              | Some m -> StringMap.add name (x, typ) m))
+          st.context.imports;
+      x
+  in
+  if not allow_tail_call then Var.Hashtbl.replace st.context.no_tail_call x ();
+  x, st
 
 let register_init_code code st =
   let st' = { var_count = 0; vars = Var.Map.empty; instrs = []; context = st.context } in
@@ -715,7 +720,7 @@ let function_body ~context ~param_names ~body =
       | Local (i, x, typ) -> local_types.(i) <- x, typ
       | Expr _ -> ())
     st.vars;
-  let body = Tail_call.f body in
+  let body = Tail_call.f ~no_tail_call:context.no_tail_call body in
   let param_count = List.length param_names in
   let locals =
     local_types

compiler/lib-wasm/code_generation.mli

@@ -42,6 +42,7 @@ type context =
   ; mutable globalized_variables : Code.Var.Set.t
   ; value_type : Wasm_ast.value_type
   ; mutable unit_name : string option
+  ; mutable no_tail_call : unit Code.Var.Hashtbl.t
   }
 
 val make_context : value_type:Wasm_ast.value_type -> context
@@ -156,7 +157,11 @@ val register_type : string -> (unit -> type_def t) -> Wasm_ast.var t
 val heap_type_sub : Wasm_ast.heap_type -> Wasm_ast.heap_type -> bool t
 
 val register_import :
-  ?import_module:string -> name:string -> Wasm_ast.import_desc -> Wasm_ast.var t
+     ?allow_tail_call:bool
+  -> ?import_module:string
+  -> name:string
+  -> Wasm_ast.import_desc
+  -> Wasm_ast.var t
 
 val register_global :
      Wasm_ast.var

compiler/lib-wasm/gc_target.ml

@@ -1297,7 +1297,13 @@ module Math = struct
     { W.params = List.init ~len:n ~f:(fun _ : W.value_type -> F64); result = [ F64 ] }
 
   let unary name x =
-    let* f = register_import ~import_module:"Math" ~name (Fun (float_func_type 1)) in
+    let* f =
+      register_import
+        ~allow_tail_call:false
+        ~import_module:"Math"
+        ~name
+        (Fun (float_func_type 1))
+    in
     let* x = x in
     return (W.Call (f, [ x ]))
 
@@ -1340,7 +1346,13 @@ module Math = struct
   let log10 f = unary "log10" f
 
   let binary name x y =
-    let* f = register_import ~import_module:"Math" ~name (Fun (float_func_type 2)) in
+    let* f =
+      register_import
+        ~allow_tail_call:false
+        ~import_module:"Math"
+        ~name
+        (Fun (float_func_type 2))
+    in
     let* x = x in
     let* y = y in
     return (W.Call (f, [ x; y ]))

compiler/lib-wasm/tail_call.ml

@@ -24,25 +24,30 @@ let get_return ~tail i =
   | Push (LocalGet y) when tail -> Some y
   | _ -> None
 
-let rewrite_tail_call ~y i =
+let rewrite_tail_call ~no_tail_call ~y i =
   match i with
-  | Wasm_ast.LocalSet (x, Call (symb, l)) when Code.Var.equal x y ->
+  | Wasm_ast.LocalSet (x, Call (symb, l))
+    when Code.Var.equal x y && not (Code.Var.Hashtbl.mem no_tail_call symb) ->
       Some (Wasm_ast.Return_call (symb, l))
   | LocalSet (x, Call_ref (ty, e, l)) when Code.Var.equal x y ->
       Some (Return_call_ref (ty, e, l))
   | _ -> None
 
-let rec instruction ~tail i =
+let rec instruction ~no_tail_call ~tail i =
   match i with
-  | Wasm_ast.Loop (ty, l) -> Wasm_ast.Loop (ty, instructions ~tail l)
-  | Block (ty, l) -> Block (ty, instructions ~tail l)
-  | If (ty, e, l1, l2) -> If (ty, e, instructions ~tail l1, instructions ~tail l2)
-  | Return (Some (Call (symb, l))) -> Return_call (symb, l)
+  | Wasm_ast.Loop (ty, l) -> Wasm_ast.Loop (ty, instructions ~no_tail_call ~tail l)
+  | Block (ty, l) -> Block (ty, instructions ~no_tail_call ~tail l)
+  | If (ty, e, l1, l2) ->
+      If (ty, e, instructions ~no_tail_call ~tail l1, instructions ~no_tail_call ~tail l2)
+  | Return (Some (Call (symb, l))) when not (Code.Var.Hashtbl.mem no_tail_call symb) ->
+      Return_call (symb, l)
   | Return (Some (Call_ref (ty, e, l))) -> Return_call_ref (ty, e, l)
-  | Push (Call (symb, l)) when tail -> Return_call (symb, l)
+  | Push (Call (symb, l)) when tail && not (Code.Var.Hashtbl.mem no_tail_call symb) ->
+      Return_call (symb, l)
   | Push (Call_ref (ty, e, l)) when tail -> Return_call_ref (ty, e, l)
   | Push (Call_ref _) -> i
-  | Drop (BlockExpr (typ, l)) -> Drop (BlockExpr (typ, instructions ~tail:false l))
+  | Drop (BlockExpr (typ, l)) ->
+      Drop (BlockExpr (typ, instructions ~no_tail_call ~tail:false l))
   | Drop _
   | LocalSet _
   | GlobalSet _
@@ -62,21 +67,28 @@ let rec instruction ~tail i =
   | Unreachable
   | Event _ -> i
 
-and instructions ~tail l =
+and instructions ~no_tail_call ~tail l =
   match l with
   | [] -> []
-  | [ i ] -> [ instruction ~tail i ]
-  | i :: Nop :: rem -> instructions ~tail (i :: rem)
-  | i :: i' :: Nop :: rem -> instructions ~tail (i :: i' :: rem)
+  | [ i ] -> [ instruction ~no_tail_call ~tail i ]
+  | i :: Nop :: rem -> instructions ~no_tail_call ~tail (i :: rem)
+  | i :: i' :: Nop :: rem -> instructions ~no_tail_call ~tail (i :: i' :: rem)
   | i :: i' :: (([] | [ Event _ ]) as event_opt) -> (
       (* There can be an event at the end of the function, which we
          should keep. *)
       match get_return ~tail i' with
-      | None -> instruction ~tail:false i :: instruction ~tail i' :: event_opt
+      | None ->
+          instruction ~no_tail_call ~tail:false i
+          :: instruction ~no_tail_call ~tail i'
+          :: event_opt
       | Some y -> (
-          match rewrite_tail_call ~y i with
-          | None -> instruction ~tail:false i :: instruction ~tail i' :: event_opt
+          match rewrite_tail_call ~no_tail_call ~y i with
+          | None ->
+              instruction ~no_tail_call ~tail:false i
+              :: instruction ~no_tail_call ~tail i'
+              :: event_opt
           | Some i'' -> i'' :: event_opt))
-  | i :: rem -> instruction ~tail:false i :: instructions ~tail rem
+  | i :: rem ->
+      instruction ~no_tail_call ~tail:false i :: instructions ~no_tail_call ~tail rem
 
-let f l = instructions ~tail:true l
+let f ~no_tail_call l = instructions ~no_tail_call ~tail:true l

compiler/lib-wasm/tail_call.mli

@@ -16,4 +16,7 @@
  * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  *)
 
-val f : Wasm_ast.instruction list -> Wasm_ast.instruction list
+val f :
+     no_tail_call:unit Code.Var.Hashtbl.t
+  -> Wasm_ast.instruction list
+  -> Wasm_ast.instruction list