inline fast path of caml_applyN (#934)

* Add compilation flag -caml-apply-inline-fast-path
* Use -caml-apply-inline-fast-path for main build and libraries
* Inline fast path of caml_apply
* Bind args
* Refactor and apply the same transformation to Flamba 2.
* Pass [ty] argument to [call_caml_apply]
* Add a CR about a version of caml_applyN that has only cold path
and update a comment.

Author: gretay-js

backend/cmm_helpers.ml

@@ -39,6 +39,14 @@ let bind_nonvar name arg fn =
     let id = V.create_local name in
     Clet (VP.create id, arg, fn (Cvar id))
 
+let bind_list name args fn =
+  let rec aux bound_args = function
+    | [] -> fn bound_args
+    | arg :: args ->
+      bind name arg (fun bound_arg -> aux (bound_arg :: bound_args) args)
+  in
+  aux [] (List.rev args)
+
 let caml_black = Nativeint.shift_left (Nativeint.of_int 3) 8
 
 let caml_local = Nativeint.shift_left (Nativeint.of_int 2) 8
@@ -2073,18 +2081,53 @@ let ptr_offset ptr offset dbg =
 let direct_apply lbl args (pos, _mode) dbg =
   Cop (Capply (typ_val, pos), Cconst_symbol (lbl, dbg) :: args, dbg)
 
+let call_caml_apply ty mut clos args pos mode dbg =
+  let arity = List.length args in
+  let really_call_caml_apply clos args =
+    let cargs =
+      (Cconst_symbol (apply_function_sym arity mode, dbg) :: args) @ [clos]
+    in
+    Cop (Capply (ty, pos), cargs, dbg)
+  in
+  if !Flambda_backend_flags.caml_apply_inline_fast_path
+  then
+    (* Generate the following expression:
+     *  (if (= clos.arity N)
+     *      (app clos.direct a1 ... aN clos)
+     *      (app caml_applyN a1 .. aN clos)
+     *)
+    (* CR-someday gyorsh: in the [else] case above, call another version of
+       caml_applyN that has only the cold path. *)
+    bind_list "arg" args (fun args ->
+        bind "fun" clos (fun clos ->
+            Cifthenelse
+              ( Cop
+                  ( Ccmpi Ceq,
+                    [ Cop
+                        ( Casr,
+                          [ get_field_gen mut clos 1 dbg;
+                            Cconst_int (pos_arity_in_closinfo, dbg) ],
+                          dbg );
+                      Cconst_int (arity, dbg) ],
+                    dbg ),
+                dbg,
+                Cop
+                  ( Capply (ty, pos),
+                    (get_field_gen mut clos 2 dbg :: args) @ [clos],
+                    dbg ),
+                dbg,
+                really_call_caml_apply clos args,
+                dbg,
+                Vval Pgenval )))
+  else really_call_caml_apply clos args
+
 let generic_apply mut clos args (pos, mode) dbg =
   match args with
   | [arg] ->
     bind "fun" clos (fun clos ->
         Cop
           (Capply (typ_val, pos), [get_field_gen mut clos 0 dbg; arg; clos], dbg))
-  | _ ->
-    let arity = List.length args in
-    let cargs =
-      (Cconst_symbol (apply_function_sym arity mode, dbg) :: args) @ [clos]
-    in
-    Cop (Capply (typ_val, pos), cargs, dbg)
+  | _ -> call_caml_apply typ_val mut clos args pos mode dbg
 
 let send kind met obj args akind dbg =
   let call_met obj args clos =
@@ -3724,12 +3767,7 @@ let indirect_call ~dbg ty pos alloc_mode f args =
            ( Capply (ty, pos),
              [load ~dbg Word_int Asttypes.Mutable ~addr:(Cvar v); arg; Cvar v],
              dbg ))
-  | args ->
-    let arity = List.length args in
-    let l =
-      (Cconst_symbol (apply_function_sym arity alloc_mode, dbg) :: args) @ [f]
-    in
-    Cop (Capply (ty, pos), l, dbg)
+  | args -> call_caml_apply ty Asttypes.Mutable f args pos alloc_mode dbg
 
 let indirect_full_call ~dbg ty pos alloc_mode f = function
   (* the single-argument case is already optimized by indirect_call *)

driver/flambda_backend_args.ml

@@ -74,6 +74,9 @@ let mk_no_long_frames f =
 let mk_debug_long_frames_threshold f =
   "-debug-long-frames-threshold", Arg.Int f, "n debug only: set long frames threshold"
 
+let mk_caml_apply_inline_fast_path f =
+  "-caml-apply-inline-fast-path", Arg.Unit f, " Inline the fast path of caml_applyN"
+
 let mk_dump_inlining_paths f =
   "-dump-inlining-paths", Arg.Unit f, " Dump inlining paths when dumping flambda2 terms"
 
@@ -476,6 +479,7 @@ module type Flambda_backend_options = sig
   val no_long_frames : unit -> unit
   val long_frames_threshold : int -> unit
 
+  val caml_apply_inline_fast_path : unit -> unit
   val internal_assembler : unit -> unit
 
   val flambda2_join_points : unit -> unit
@@ -557,6 +561,8 @@ struct
     mk_no_long_frames F.no_long_frames;
     mk_debug_long_frames_threshold F.long_frames_threshold;
 
+    mk_caml_apply_inline_fast_path F.caml_apply_inline_fast_path;
+
     mk_internal_assembler F.internal_assembler;
 
     mk_flambda2_join_points F.flambda2_join_points;
@@ -674,6 +680,9 @@ module Flambda_backend_options_impl = struct
   let no_long_frames = clear' Flambda_backend_flags.allow_long_frames
   let long_frames_threshold n = set_long_frames_threshold n
 
+  let caml_apply_inline_fast_path =
+    set' Flambda_backend_flags.caml_apply_inline_fast_path
+
   let internal_assembler = set' Flambda_backend_flags.internal_assembler
 
   let flambda2_join_points = set Flambda2.join_points
@@ -882,6 +891,8 @@ module Extra_params = struct
              (Printf.sprintf "Expected integer between 0 and %d"
                 Flambda_backend_flags.max_long_frames_threshold))
       end
+    | "caml-apply-inline-fast-path" ->
+      set' Flambda_backend_flags.caml_apply_inline_fast_path
     | "dasm-comments" -> set' Flambda_backend_flags.dasm_comments
     | "gupstream-dwarf" -> set' Debugging.restrict_to_upstream_dwarf
     | "gstartup" -> set' Debugging.dwarf_for_startup_file

driver/flambda_backend_args.mli

@@ -43,6 +43,7 @@ module type Flambda_backend_options = sig
   val no_long_frames : unit -> unit
   val long_frames_threshold : int -> unit
 
+  val caml_apply_inline_fast_path : unit -> unit
   val internal_assembler : unit -> unit
 
   val flambda2_join_points : unit -> unit

driver/flambda_backend_flags.ml

@@ -35,6 +35,8 @@ let allow_long_frames = ref true        (* -no-long-frames *)
 let max_long_frames_threshold = 0x7FFF
 let long_frames_threshold = ref max_long_frames_threshold (* -debug-long-frames-threshold n *)
 
+let caml_apply_inline_fast_path = ref false  (* -caml-apply-inline-fast-path *)
+
 type function_result_types = Never | Functors_only | All_functions
 type opt_level = Oclassic | O2 | O3
 type 'a or_default = Set of 'a | Default

driver/flambda_backend_flags.mli

@@ -32,6 +32,7 @@ val disable_poll_insertion : bool ref
 val allow_long_frames : bool ref
 val max_long_frames_threshold : int
 val long_frames_threshold : int ref
+val caml_apply_inline_fast_path : bool ref
 
 type function_result_types = Never | Functors_only | All_functions
 type opt_level = Oclassic | O2 | O3

ocaml/dune

@@ -25,7 +25,7 @@
    ;; CR gyorsh: it is not clear what the ":standard" flags are, and they
    ;; may change depending on the version of dune.
    ;; Consider hard-coded flags, such as -O3.
-   (:standard -alloc-check)))
+   (:standard -alloc-check -caml-apply-inline-fast-path)))
  (boot
   (flags
    (:standard -warn-error +A))))