flambda2-types: New n-way join algorithm

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 ocaml-flambda#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 ocaml-flambda#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

driver/flambda_backend_flags.ml

@@ -134,7 +134,7 @@ module Flambda2 = struct
     let cse_depth = 2
     let join_depth = 5
     let function_result_types = Never
-    let meet_algorithm = Basic
+    let meet_algorithm = Advanced
     let enable_reaper = false
     let unicode = true
   end

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
@@ -274,4 +421,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 ()