Do not add extra_params to the lifted_cont_params (#3885)

Currently, any and all variables in scope, including parameters
and extra parameters (coming from CSE and unboxing), are added to the
lifted_cont_params. This is superfluous, since the handlers of lifted
continuations do not contain any reference to these extra params prior
to simplification, and these extra params can be recomputed for the
lifted continuations.

As it stands, this commit alone may lead to regressions, since we
currently do not perform unboxing on the lifted_cont_params. This will
be fixed in a later commit.

Author: Gbury

middle_end/flambda2/simplify/env/downwards_env.ml

@@ -116,9 +116,9 @@ let [@ocamlformat "disable"] print ppf { round; typing_env;
     (Format.pp_print_list ~pp_sep:Format.pp_print_space Lifted_cont_params.print) defined_variables_by_scope
     cost_of_lifting_continuations_out_of_current_one
 
-let define_variable t var kind =
+let define_variable0 ~extra t var kind =
   let defined_variables_by_scope =
-    if Variable.Set.mem (Bound_var.var var) t.lifted
+    if extra || Variable.Set.mem (Bound_var.var var) t.lifted
     then t.defined_variables_by_scope
     else
       match t.defined_variables_by_scope with
@@ -146,6 +146,9 @@ let define_variable t var kind =
     defined_variables_by_scope
   }
 
+let define_variable t var kind =
+  (define_variable0 [@inlined hint]) ~extra:false t var kind
+
 let create ~round ~(resolver : resolver)
     ~(get_imported_names : get_imported_names)
     ~(get_imported_code : get_imported_code) ~propagating_float_consts
@@ -290,12 +293,15 @@ let define_name t name kind =
         kind)
     ~symbol:(fun [@inline] sym -> (define_symbol [@inlined hint]) t sym kind)
 
-let add_variable t var ty =
-  let t = (define_variable [@inlined hint]) t var (T.kind ty) in
+let add_variable0 ~extra t var ty =
+  let t = (define_variable0 [@inlined hint]) ~extra t var (T.kind ty) in
   { t with
     typing_env = TE.add_equation t.typing_env (Name.var (Bound_var.var var)) ty
   }
 
+let add_variable t var ty =
+  (add_variable0 [@inlined hint]) ~extra:false t var ty
+
 let add_symbol t sym ty =
   let t = (define_symbol [@inlined hint]) t sym (T.kind ty) in
   { t with typing_env = TE.add_equation t.typing_env (Name.symbol sym) ty }
@@ -339,15 +345,16 @@ let add_equation_on_name t name ty =
   let typing_env = TE.add_equation t.typing_env name ty in
   { t with typing_env }
 
-let define_parameters t ~params =
+let define_parameters ~extra t ~params =
   List.fold_left
     (fun t param ->
       let var = Bound_var.create (BP.var param) Name_mode.normal in
-      define_variable t var (K.With_subkind.kind (BP.kind param)))
+      define_variable0 ~extra t var (K.With_subkind.kind (BP.kind param)))
     t
     (Bound_parameters.to_list params)
 
-let add_parameters ?(name_mode = Name_mode.normal) t params ~param_types =
+let add_parameters ~extra ?(name_mode = Name_mode.normal) t params ~param_types
+    =
   let params' = params in
   let params = Bound_parameters.to_list params in
   if List.compare_lengths params param_types <> 0
@@ -360,10 +367,10 @@ let add_parameters ?(name_mode = Name_mode.normal) t params ~param_types =
   List.fold_left2
     (fun t param param_type ->
       let var = Bound_var.create (BP.var param) name_mode in
-      add_variable t var param_type)
+      add_variable0 ~extra t var param_type)
     t params param_types
 
-let add_parameters_with_unknown_types ?alloc_modes ?name_mode t params =
+let add_parameters_with_unknown_types ~extra ?alloc_modes ?name_mode t params =
   let params' = params in
   let params = Bound_parameters.to_list params in
   let alloc_modes =
@@ -380,7 +387,7 @@ let add_parameters_with_unknown_types ?alloc_modes ?name_mode t params =
     ListLabels.map2 params alloc_modes ~f:(fun param alloc_mode ->
         T.unknown_with_subkind ~alloc_mode (BP.kind param))
   in
-  add_parameters ?name_mode t params' ~param_types
+  add_parameters ~extra ?name_mode t params' ~param_types
 
 let mark_parameters_as_toplevel t params =
   let variables_defined_at_toplevel =

middle_end/flambda2/simplify/env/downwards_env.mli

@@ -99,16 +99,18 @@ val define_name_if_undefined : t -> Bound_name.t -> Flambda_kind.t -> t
 
 val add_equation_on_name : t -> Name.t -> Flambda2_types.t -> t
 
-val define_parameters : t -> params:Bound_parameters.t -> t
+val define_parameters : extra:bool -> t -> params:Bound_parameters.t -> t
 
 val add_parameters :
+  extra:bool ->
   ?name_mode:Name_mode.t ->
   t ->
   Bound_parameters.t ->
   param_types:Flambda2_types.t list ->
   t
 
 val add_parameters_with_unknown_types :
+  extra:bool ->
   ?alloc_modes:Alloc_mode.For_types.t list ->
   ?name_mode:Name_mode.t ->
   t ->

middle_end/flambda2/simplify/join_points.ml

@@ -65,7 +65,7 @@ let introduce_extra_params_for_join denv use_envs_with_ids
     denv, use_tenvs_with_ids
   else
     let extra_params = EPA.extra_params extra_params_and_args in
-    let denv = DE.define_parameters denv ~params:extra_params in
+    let denv = DE.define_parameters ~extra:true denv ~params:extra_params in
     let use_envs_with_ids =
       List.filter_map
         (introduce_extra_params_in_use_env extra_params_and_args)
@@ -185,7 +185,9 @@ let compute_handler_env ?cut_after uses ~is_recursive ~env_at_fork
           add_equations_on_params typing_env ~is_recursive ~params ~param_types
         in
         let use_env =
-          let use_env = DE.define_parameters (U.env_at_use use) ~params in
+          let use_env =
+            DE.define_parameters ~extra:false (U.env_at_use use) ~params
+          in
           DE.map_typing_env use_env ~f:add_or_meet_param_type
         in
         use_env, U.id use, U.use_kind use)
@@ -266,15 +268,17 @@ let compute_handler_env ?cut_after uses ~is_recursive ~env_at_fork
       then
         (* No need to add equations, as they will be computed from the use
            environments *)
-        let denv = DE.define_parameters denv ~params in
+        let denv = DE.define_parameters denv ~extra:false ~params in
         Profile.record_call ~accumulate:true "join" (fun () ->
             join ?cut_after denv params ~consts_lifted_during_body
               ~use_envs_with_ids ~lifted_cont_extra_params_and_args)
       else
         (* Define parameters with basic equations from the subkinds *)
-        let denv = DE.add_parameters_with_unknown_types denv params in
         let denv =
-          DE.add_parameters_with_unknown_types denv
+          DE.add_parameters_with_unknown_types denv ~extra:false params
+        in
+        let denv =
+          DE.add_parameters_with_unknown_types denv ~extra:true
             (EPA.extra_params lifted_cont_extra_params_and_args)
         in
         denv, lifted_cont_extra_params_and_args

middle_end/flambda2/simplify/simplify_apply_expr.ml

@@ -296,7 +296,7 @@ let simplify_direct_full_application ~simplify_expr dacc apply function_type
                   result_types Apply.print apply;
               let denv = DA.denv dacc in
               let denv =
-                DE.add_parameters_with_unknown_types
+                DE.add_parameters_with_unknown_types ~extra:false
                   ~name_mode:Name_mode.in_types denv params
               in
               let params = Bound_parameters.to_list params in

middle_end/flambda2/simplify/simplify_let_cont_expr.ml

@@ -1137,7 +1137,7 @@ and simplify_single_recursive_handler ~simplify_expr cont_uses_env_so_far
      uses (needed for the recursive continuation handlers). *)
   let handler_env =
     assert (not (DE.at_unit_toplevel denv_to_reset));
-    DE.add_parameters_with_unknown_types denv_to_reset params
+    DE.add_parameters_with_unknown_types ~extra:false denv_to_reset params
   in
   let handler_env = LCS.add_to_denv handler_env consts_lifted_during_body in
   let code_age_relation = TE.code_age_relation (DA.typing_env dacc) in

middle_end/flambda2/simplify/simplify_set_of_closures.ml

@@ -34,7 +34,8 @@ let dacc_inside_function context ~outer_dacc ~params ~my_closure ~my_region
   let dacc = C.dacc_inside_functions context in
   let alloc_modes = Code_metadata.param_modes code_metadata in
   let denv =
-    DE.add_parameters_with_unknown_types ~alloc_modes (DA.denv dacc) params
+    DE.add_parameters_with_unknown_types ~extra:false ~alloc_modes
+      (DA.denv dacc) params
     |> DE.set_inlining_arguments inlining_arguments
     |> DE.set_inlining_history_tracker
          (Inlining_history.Tracker.inside_function absolute_history)