flambda2-types: New n-way join algorithm (#3538)

The existing join algorithm suffers from several drawbacks:

 - It can be slow due to the use of a quadratic algorithm, taking up to
   60% of the total compilation time in -O3 mode in pathological cases
   (lambda_to_flambda_primitives.ml). See also #3300.

 - It is inefficient as it computes the join of all types appearing in
   *any* joined environment prior to filtering out the types that are
   not needed, instead of first computing the types whose join will be
   needed.

 - It is sensitive to the names of local variables that only exist in
   some of the joined environments but not in the target environment.

 - It relies on a global binding time of variables across all joined
   environments and the target environment that does not exist, as
   figured in #3278. Subsequently, it can lose aliasing information,
   and breaks typing env invariants by recording the same variable as
   defined multiple times (with dubious semantics).

This patch implements a new join algorithm, based on a n-way join of types.
The new algorithm is:

 - Faster, as it avoids quadratic complexity (outside of complex nesting
   of env extensions). Compared to the existing join algorithm (with
   advanced meet), on my machine, the new join algorithm is 30x faster
   on the pathological lambda_to_flambda_primitives.ml, taking only
   around 10% of the total compilation time and speeding up the
   compilation of the file by 3.5x. On camlinternalFormat.ml, the new
   join is about 2.5-3x faster, reducing the time spent in the join from
   20% to less than 10% and speeding up the total compilation time by
   about 20%.

 - More efficient, as it only computes a join if it can possibly result
   in a more precise type, i.e. if the variable has been assigned a new
   type in all joined environments (otherwise the existing type in the
   target environment is already the most precise).

 - Independent of the names of local variables.

 - Only depends on a consistent binding time *order* of the shared
   variables (defined in both the target environment and all joined
   environments), which is respected. Since the result is independent of
   the binding times of local / existential variables, the typing env
   invariants are respected.

Author: bclement-ocp

.github/workflows/build.yml

@@ -58,14 +58,14 @@ jobs:
             config: --enable-middle-end=flambda2 --enable-frame-pointers --enable-runtime5 --enable-poll-insertion --enable-flambda-invariants
             os: ubuntu-latest
             build_ocamlparam: ''
-            ocamlparam: '_,O3=1,flambda2-expert-cont-lifting-budget=200'
+            ocamlparam: '_,O3=1,flambda2-expert-cont-lifting-budget=200,flambda2-join-algorithm=n-way'
 
           - name: flambda2_o3_advanced_meet_frame_pointers_runtime5_debug
             config: --enable-middle-end=flambda2 --enable-frame-pointers --enable-runtime5
             os: ubuntu-latest
             build_ocamlparam: ''
             use_runtime: d
-            ocamlparam: '_,O3=1,flambda2-expert-cont-lifting-budget=200,cfg-invariants=1,cfg-eliminate-dead-trap-handlers=1'
+            ocamlparam: '_,O3=1,flambda2-expert-cont-lifting-budget=200,cfg-invariants=1,cfg-eliminate-dead-trap-handlers=1,flambda2-join-algorithm=n-way'
 
           - name: flambda2_frame_pointers_oclassic_polling
             config: --enable-middle-end=flambda2 --enable-frame-pointers --enable-poll-insertion --enable-flambda-invariants

driver/flambda_backend_args.ml

@@ -268,6 +268,15 @@ let mk_flambda2_advanced_meet f =
   Printf.sprintf " Use an advanced meet algorithm (deprecated) (Flambda 2 only)"
 ;;
 
+let mk_flambda2_join_algorithm f =
+  "-flambda2-join-algorithm", Arg.Symbol (["binary"; "n-way"; "checked"], f),
+  Printf.sprintf " Select the join algorithm to use (Flambda 2 only)\n \
+      \     Valid values are: \n\
+      \       \"binary\" is the legacy binary join;\n\
+      \       \"n-way\" is the new n-way join;\n\
+      \       \"checked\" runs both algorithms and compares them (use for \
+      debugging)."
+;;
 
 let mk_flambda2_join_points f =
   "-flambda2-join-points", Arg.Unit f,
@@ -777,6 +786,7 @@ module type Flambda_backend_options = sig
   val no_flambda2_result_types : unit -> unit
   val flambda2_basic_meet : unit -> unit
   val flambda2_advanced_meet : unit -> unit
+  val flambda2_join_algorithm : string -> unit
   val flambda2_unbox_along_intra_function_control_flow : unit -> unit
   val no_flambda2_unbox_along_intra_function_control_flow : unit -> unit
   val flambda2_backend_cse_at_toplevel : unit -> unit
@@ -916,6 +926,7 @@ struct
       F.no_flambda2_result_types;
     mk_flambda2_basic_meet F.flambda2_basic_meet;
     mk_flambda2_advanced_meet F.flambda2_advanced_meet;
+    mk_flambda2_join_algorithm F.flambda2_join_algorithm;
     mk_flambda2_unbox_along_intra_function_control_flow
       F.flambda2_unbox_along_intra_function_control_flow;
     mk_no_flambda2_unbox_along_intra_function_control_flow
@@ -1126,6 +1137,15 @@ module Flambda_backend_options_impl = struct
     Flambda2.function_result_types := Flambda_backend_flags.Set Flambda_backend_flags.Never
   let flambda2_basic_meet () = ()
   let flambda2_advanced_meet () = ()
+  let flambda2_join_algorithm algorithm =
+    match algorithm with
+    | "binary" ->
+      Flambda2.join_algorithm := Flambda_backend_flags.Set Flambda_backend_flags.Binary
+    | "n-way" ->
+      Flambda2.join_algorithm := Flambda_backend_flags.Set Flambda_backend_flags.N_way
+    | "checked" ->
+      Flambda2.join_algorithm := Flambda_backend_flags.Set Flambda_backend_flags.Checked
+    | _ -> () (* This should not occur as we use Arg.Symbol *)
   let flambda2_unbox_along_intra_function_control_flow =
     set Flambda2.unbox_along_intra_function_control_flow
   let no_flambda2_unbox_along_intra_function_control_flow =
@@ -1456,6 +1476,13 @@ module Extra_params = struct
       | _ ->
         Misc.fatal_error "Syntax: flambda2-meet_algorithm=basic|advanced");
       true
+    | "flambda2-join-algorithm" ->
+      (match String.lowercase_ascii v with
+      | "binary" | "n-way" | "checked" as v ->
+        Flambda_backend_options_impl.flambda2_join_algorithm v
+      | _ ->
+        Misc.fatal_error "Syntax: flambda2-join-algorithm=binary|n-way|checked");
+      true
     | "flambda2-unbox-along-intra-function-control-flow" ->
        set Flambda2.unbox_along_intra_function_control_flow
     | "flambda2-backend-cse-at-toplevel" ->

driver/flambda_backend_args.mli

@@ -91,6 +91,7 @@ module type Flambda_backend_options = sig
   val no_flambda2_result_types : unit -> unit
   val flambda2_basic_meet : unit -> unit
   val flambda2_advanced_meet : unit -> unit
+  val flambda2_join_algorithm : string -> unit
   val flambda2_unbox_along_intra_function_control_flow : unit -> unit
   val no_flambda2_unbox_along_intra_function_control_flow : unit -> unit
   val flambda2_backend_cse_at_toplevel : unit -> unit

driver/flambda_backend_flags.ml

@@ -96,6 +96,7 @@ let long_frames_threshold = ref max_long_frames_threshold (* -debug-long-frames-
 let caml_apply_inline_fast_path = ref false  (* -caml-apply-inline-fast-path *)
 
 type function_result_types = Never | Functors_only | All_functions
+type join_algorithm = Binary | N_way | Checked
 type opt_level = Oclassic | O2 | O3
 type 'a or_default = Set of 'a | Default
 
@@ -128,6 +129,7 @@ module Flambda2 = struct
     let backend_cse_at_toplevel = false
     let cse_depth = 2
     let join_depth = 5
+    let join_algorithm = Binary
     let function_result_types = Never
     let enable_reaper = false
     let unicode = true
@@ -141,6 +143,7 @@ module Flambda2 = struct
     backend_cse_at_toplevel : bool;
     cse_depth : int;
     join_depth : int;
+    join_algorithm : join_algorithm;
     function_result_types : function_result_types;
     enable_reaper : bool;
     unicode : bool;
@@ -154,6 +157,7 @@ module Flambda2 = struct
     backend_cse_at_toplevel = Default.backend_cse_at_toplevel;
     cse_depth = Default.cse_depth;
     join_depth = Default.join_depth;
+    join_algorithm = Default.join_algorithm;
     function_result_types = Default.function_result_types;
     enable_reaper = Default.enable_reaper;
     unicode = Default.unicode;
@@ -187,6 +191,7 @@ module Flambda2 = struct
   let backend_cse_at_toplevel = ref Default
   let cse_depth = ref Default
   let join_depth = ref Default
+  let join_algorithm = ref Default
   let unicode = ref Default
   let kind_checks = ref Default
   let function_result_types = ref Default

driver/flambda_backend_flags.mli

@@ -83,6 +83,7 @@ val long_frames_threshold : int ref
 val caml_apply_inline_fast_path : bool ref
 
 type function_result_types = Never | Functors_only | All_functions
+type join_algorithm = Binary | N_way | Checked
 type opt_level = Oclassic | O2 | O3
 type 'a or_default = Set of 'a | Default
 
@@ -109,6 +110,7 @@ module Flambda2 : sig
     val backend_cse_at_toplevel : bool
     val cse_depth : int
     val join_depth : int
+    val join_algorithm : join_algorithm
     val function_result_types : function_result_types
     val enable_reaper : bool
     val unicode : bool
@@ -125,6 +127,7 @@ module Flambda2 : sig
     backend_cse_at_toplevel : bool;
     cse_depth : int;
     join_depth : int;
+    join_algorithm : join_algorithm;
     function_result_types : function_result_types;
     enable_reaper : bool;
     unicode : bool;
@@ -141,6 +144,7 @@ module Flambda2 : sig
   val backend_cse_at_toplevel : bool or_default ref
   val cse_depth : int or_default ref
   val join_depth : int or_default ref
+  val join_algorithm : join_algorithm or_default ref
   val enable_reaper : bool or_default ref
   val unicode : bool or_default ref
   val kind_checks : bool or_default ref

middle_end/flambda2/tests/meet_test.ml

@@ -139,6 +139,153 @@ let meet_variants_don't_lose_aliases () =
       Format.eprintf "@[<hov 2>meet:@ %a@]@.@[<hov 2>env:@ %a@]@." T.print
         tag_meet_ty TE.print tag_meet_env)
 
+let test_join_with_extensions () =
+  let define ?(kind = K.value) env v =
+    let v' = Bound_var.create v Name_mode.normal in
+    TE.add_definition env (Bound_name.create_var v') kind
+  in
+  let env = create_env () in
+  let y = Variable.create "y" in
+  let x = Variable.create "x" in
+  let a = Variable.create "a" in
+  let b = Variable.create "b" in
+  let env = define env y in
+  let env = define env x in
+  let env = define ~kind:K.naked_immediate env a in
+  let env = define ~kind:K.naked_immediate env b in
+  let tag_0 = Tag.Scannable.zero in
+  let tag_1 = Option.get (Tag.Scannable.of_tag (Tag.create_exn 1)) in
+  let make ty =
+    T.variant
+      ~const_ctors:(T.bottom K.naked_immediate)
+      ~non_const_ctors:
+        (Tag.Scannable.Map.of_list
+           [ tag_0, (K.Block_shape.Scannable Value_only, [ty]);
+             tag_1, (K.Block_shape.Scannable Value_only, []) ])
+      Alloc_mode.For_types.heap
+  in
+  let env = TE.add_equation env (Name.var y) (make (T.unknown K.value)) in
+  let scope = TE.current_scope env in
+  let scoped_env = TE.increment_scope env in
+  let left_env =
+    TE.add_equation scoped_env (Name.var x)
+      (T.tagged_immediate_alias_to ~naked_immediate:a)
+  in
+  let right_env =
+    TE.add_equation scoped_env (Name.var x)
+      (T.tagged_immediate_alias_to ~naked_immediate:b)
+  in
+  let ty_a = make (T.tagged_immediate_alias_to ~naked_immediate:a) in
+  let ty_b = make (T.tagged_immediate_alias_to ~naked_immediate:b) in
+  let left_env = TE.add_equation left_env (Name.var y) ty_a in
+  let right_env =
+    match T.meet right_env ty_a ty_b with
+    | Ok (ty, right_env) -> TE.add_equation right_env (Name.var y) ty
+    | Bottom -> assert false
+  in
+  Format.eprintf "Left:@.%a@." TE.print left_env;
+  Format.eprintf "Right:@.%a@." TE.print right_env;
+  let joined_env =
+    T.cut_and_n_way_join scoped_env
+      [ left_env, Apply_cont_rewrite_id.create (), Inlinable;
+        right_env, Apply_cont_rewrite_id.create (), Inlinable ]
+      ~params:Bound_parameters.empty ~cut_after:scope
+      ~extra_allowed_names:Name_occurrences.empty
+      ~extra_lifted_consts_in_use_envs:Symbol.Set.empty
+  in
+  Format.eprintf "Res:@.%a@." TE.print joined_env
+
+let test_join_with_complex_extensions () =
+  let define ?(kind = K.value) env v =
+    let v' = Bound_var.create v Name_mode.normal in
+    TE.add_definition env (Bound_name.create_var v') kind
+  in
+  let env = create_env () in
+  let y = Variable.create "y" in
+  let x = Variable.create "x" in
+  let w = Variable.create "w" in
+  let z = Variable.create "z" in
+  let a = Variable.create "a" in
+  let b = Variable.create "b" in
+  let c = Variable.create "c" in
+  let d = Variable.create "d" in
+  let env = define env z in
+  let env = define env x in
+  let env = define env y in
+  let env = define env w in
+  let env = define ~kind:K.naked_immediate env a in
+  let env = define ~kind:K.naked_immediate env b in
+  let env = define ~kind:K.naked_immediate env c in
+  let env = define ~kind:K.naked_immediate env d in
+  let tag_0 = Tag.Scannable.zero in
+  let tag_1 = Option.get (Tag.Scannable.of_tag (Tag.create_exn 1)) in
+  let make tys =
+    T.variant
+      ~const_ctors:(T.bottom K.naked_immediate)
+      ~non_const_ctors:
+        (Tag.Scannable.Map.of_list
+           [ tag_0, (K.Block_shape.Scannable Value_only, tys);
+             tag_1, (K.Block_shape.Scannable Value_only, []) ])
+      Alloc_mode.For_types.heap
+  in
+  let env =
+    TE.add_equation env (Name.var z)
+      (make [T.unknown K.value; T.unknown K.value])
+  in
+  let scope = TE.current_scope env in
+  let scoped_env = TE.increment_scope env in
+  let left_env =
+    TE.add_equation scoped_env (Name.var x)
+      (T.tagged_immediate_alias_to ~naked_immediate:a)
+  in
+  let left_env =
+    TE.add_equation left_env (Name.var y)
+      (T.tagged_immediate_alias_to ~naked_immediate:a)
+  in
+  let left_env =
+    TE.add_equation left_env (Name.var w)
+      (T.tagged_immediate_alias_to ~naked_immediate:a)
+  in
+  let right_env =
+    TE.add_equation scoped_env (Name.var x)
+      (T.tagged_immediate_alias_to ~naked_immediate:b)
+  in
+  let right_env =
+    TE.add_equation right_env (Name.var y)
+      (T.tagged_immediate_alias_to ~naked_immediate:c)
+  in
+  let right_env =
+    TE.add_equation right_env (Name.var w)
+      (T.tagged_immediate_alias_to ~naked_immediate:d)
+  in
+  let ty_a =
+    make
+      [ T.tagged_immediate_alias_to ~naked_immediate:b;
+        T.tagged_immediate_alias_to ~naked_immediate:b ]
+  in
+  let ty_b =
+    make
+      [ T.tagged_immediate_alias_to ~naked_immediate:c;
+        T.tagged_immediate_alias_to ~naked_immediate:d ]
+  in
+  let left_env = TE.add_equation left_env (Name.var z) ty_a in
+  let right_env =
+    match T.meet right_env ty_a ty_b with
+    | Ok (ty, right_env) -> TE.add_equation right_env (Name.var z) ty
+    | Bottom -> assert false
+  in
+  Format.eprintf "Left:@.%a@." TE.print left_env;
+  Format.eprintf "Right:@.%a@." TE.print right_env;
+  let joined_env =
+    T.cut_and_n_way_join scoped_env
+      [ left_env, Apply_cont_rewrite_id.create (), Inlinable;
+        right_env, Apply_cont_rewrite_id.create (), Inlinable ]
+      ~params:Bound_parameters.empty ~cut_after:scope
+      ~extra_allowed_names:Name_occurrences.empty
+      ~extra_lifted_consts_in_use_envs:Symbol.Set.empty
+  in
+  Format.eprintf "Res:@.%a@." TE.print joined_env
+
 let test_meet_two_blocks () =
   let define env v =
     let v' = Bound_var.create v Name_mode.normal in
@@ -272,4 +419,8 @@ let () =
   Format.eprintf "@.MEET ALIAS TO RECOVER @\n@.";
   test_meet_recover_alias ();
   Format.eprintf "@.MEET BOTTOM AFTER ALIAS@\n@.";
-  test_meet_bottom_after_alias ()
+  test_meet_bottom_after_alias ();
+  Format.eprintf "@.JOIN WITH EXTENSIONS@\n@.";
+  test_join_with_extensions ();
+  Format.eprintf "@.JOIN WITH COMPLEX EXTENSIONS@\n@.";
+  test_join_with_complex_extensions ()