Delete unused local allocation regions in Flambda 2 (#809)

Author: mshinwell

middle_end/flambda2/bound_identifiers/bound_for_function.ml

@@ -19,37 +19,46 @@ type t =
     exn_continuation : Continuation.t;
     params : Bound_parameters.t;
     my_closure : Variable.t;
+    my_region : Variable.t;
     my_depth : Variable.t
   }
 
 let[@ocamlformat "disable"] print ppf
-    { return_continuation; exn_continuation; params; my_closure; my_depth } =
+    { return_continuation; exn_continuation; params; my_closure; my_region; my_depth } =
   Format.fprintf ppf "@[<hov 1>(\
       @[<hov 1>(return_continuation@ %a)@]@ \
       @[<hov 1>(exn_continuation@ %a)@]@ \
       @[<hov 1>(params@ %a)@]@ \
       @[<hov 1>(my_closure@ %a)@]@ \
+      @[<hov 1>(my_region@ %a)@]@ \
       @[<hov 1>(my_depth@ %a)@])@]"
     Continuation.print return_continuation
     Continuation.print exn_continuation
     Bound_parameters.print params
     Variable.print my_closure
+    Variable.print my_region
     Variable.print my_depth
 
-let create ~return_continuation ~exn_continuation ~params ~my_closure ~my_depth
-    =
+let create ~return_continuation ~exn_continuation ~params ~my_closure ~my_region
+    ~my_depth =
   Bound_parameters.check_no_duplicates params;
   (if Flambda_features.check_invariants ()
   then
     let params_set = Bound_parameters.var_set params in
-    if Variable.equal my_closure my_depth
-       || Variable.Set.mem my_closure params_set
-       || Variable.Set.mem my_depth params_set
+    let my_set = Variable.Set.of_list [my_closure; my_region; my_depth] in
+    if Variable.Set.cardinal my_set <> 3
+       || not (Variable.Set.is_empty (Variable.Set.inter my_set params_set))
     then
       Misc.fatal_errorf
-        "[my_closure] and [my_depth] must be disjoint from themselves and the \
-         other parameters");
-  { return_continuation; exn_continuation; params; my_closure; my_depth }
+        "[my_closure], [my_region] and [my_depth] must be disjoint from \
+         themselves and the other parameters");
+  { return_continuation;
+    exn_continuation;
+    params;
+    my_closure;
+    my_region;
+    my_depth
+  }
 
 let return_continuation t = t.return_continuation
 
@@ -59,10 +68,18 @@ let params t = t.params
 
 let my_closure t = t.my_closure
 
+let my_region t = t.my_region
+
 let my_depth t = t.my_depth
 
 let free_names
-    { return_continuation; exn_continuation; params; my_closure; my_depth } =
+    { return_continuation;
+      exn_continuation;
+      params;
+      my_closure;
+      my_region;
+      my_depth
+    } =
   (* See [bound_continuations.ml] for why [add_traps] is [true]. *)
   let free_names =
     Name_occurrences.add_continuation Name_occurrences.empty return_continuation
@@ -78,11 +95,19 @@ let free_names
   let free_names =
     Name_occurrences.add_variable free_names my_closure Name_mode.normal
   in
+  let free_names =
+    Name_occurrences.add_variable free_names my_region Name_mode.normal
+  in
   Name_occurrences.add_variable free_names my_depth Name_mode.normal
 
 let apply_renaming
-    { return_continuation; exn_continuation; params; my_closure; my_depth }
-    renaming =
+    { return_continuation;
+      exn_continuation;
+      params;
+      my_closure;
+      my_region;
+      my_depth
+    } renaming =
   let return_continuation =
     Renaming.apply_continuation renaming return_continuation
   in
@@ -91,25 +116,47 @@ let apply_renaming
   in
   let params = Bound_parameters.apply_renaming params renaming in
   let my_closure = Renaming.apply_variable renaming my_closure in
+  let my_region = Renaming.apply_variable renaming my_region in
   let my_depth = Renaming.apply_variable renaming my_depth in
-  { return_continuation; exn_continuation; params; my_closure; my_depth }
+  (* CR mshinwell: this should have a phys-equal check *)
+  { return_continuation;
+    exn_continuation;
+    params;
+    my_closure;
+    my_region;
+    my_depth
+  }
 
 let ids_for_export
-    { return_continuation; exn_continuation; params; my_closure; my_depth } =
+    { return_continuation;
+      exn_continuation;
+      params;
+      my_closure;
+      my_region;
+      my_depth
+    } =
   let ids =
     Ids_for_export.add_continuation Ids_for_export.empty return_continuation
   in
   let ids = Ids_for_export.add_continuation ids exn_continuation in
   let ids = Ids_for_export.union ids (Bound_parameters.ids_for_export params) in
   let ids = Ids_for_export.add_variable ids my_closure in
+  let ids = Ids_for_export.add_variable ids my_region in
   Ids_for_export.add_variable ids my_depth
 
 let rename
-    { return_continuation; exn_continuation; params; my_closure; my_depth } =
+    { return_continuation;
+      exn_continuation;
+      params;
+      my_closure;
+      my_region;
+      my_depth
+    } =
   { return_continuation = Continuation.rename return_continuation;
     exn_continuation = Continuation.rename exn_continuation;
     params = Bound_parameters.rename params;
     my_closure = Variable.rename my_closure;
+    my_region = Variable.rename my_region;
     my_depth = Variable.rename my_depth
   }
 
@@ -118,13 +165,15 @@ let renaming
       exn_continuation = exn_continuation1;
       params = params1;
       my_closure = my_closure1;
+      my_region = my_region1;
       my_depth = my_depth1
     }
     ~guaranteed_fresh:
       { return_continuation = return_continuation2;
         exn_continuation = exn_continuation2;
         params = params2;
         my_closure = my_closure2;
+        my_region = my_region2;
         my_depth = my_depth2
       } =
   let renaming =
@@ -144,4 +193,7 @@ let renaming
     Renaming.add_fresh_variable renaming my_closure1
       ~guaranteed_fresh:my_closure2
   in
+  let renaming =
+    Renaming.add_fresh_variable renaming my_region1 ~guaranteed_fresh:my_region2
+  in
   Renaming.add_fresh_variable renaming my_depth1 ~guaranteed_fresh:my_depth2

middle_end/flambda2/bound_identifiers/bound_for_function.mli

@@ -24,6 +24,7 @@ val create :
   exn_continuation:Continuation.t ->
   params:Bound_parameters.t ->
   my_closure:Variable.t ->
+  my_region:Variable.t ->
   my_depth:Variable.t ->
   t
 
@@ -35,6 +36,8 @@ val params : t -> Bound_parameters.t
 
 val my_closure : t -> Variable.t
 
+val my_region : t -> Variable.t
+
 val my_depth : t -> Variable.t
 
 include Bindable.S with type t := t

middle_end/flambda2/compare/compare.ml

@@ -274,7 +274,7 @@ let subst_set_of_closures env set =
            subst_value_slot env var, subst_simple env simple)
     |> Value_slot.Map.of_list
   in
-  Set_of_closures.create Heap ~value_slots decls
+  Set_of_closures.create Alloc_mode.With_region.heap ~value_slots decls
 
 let subst_rec_info_expr _env ri =
   (* Only depth variables can occur in [Rec_info_expr], and we only mess with
@@ -290,7 +290,7 @@ let subst_call_kind env (call_kind : Call_kind.t) : Call_kind.t =
   match call_kind with
   | Function { function_call = Direct { code_id; return_arity }; _ } ->
     let code_id = subst_code_id env code_id in
-    Call_kind.direct_function_call code_id ~return_arity Heap
+    Call_kind.direct_function_call code_id ~return_arity Alloc_mode.heap
   | _ -> call_kind
 
 let rec subst_expr env e =
@@ -387,12 +387,13 @@ and subst_params_and_body env params_and_body =
          ~body
          ~my_closure
          ~is_my_closure_used:_
+         ~my_region
          ~my_depth
          ~free_names_of_body
        ->
       let body = subst_expr env body in
       Function_params_and_body.create ~return_continuation ~exn_continuation
-        params ~body ~my_closure ~free_names_of_body ~my_depth)
+        params ~body ~my_closure ~my_region ~free_names_of_body ~my_depth)
 
 and subst_let_cont env (let_cont_expr : Let_cont_expr.t) =
   match let_cont_expr with
@@ -432,8 +433,10 @@ and subst_apply env apply =
   let inlining_state = Apply_expr.inlining_state apply in
   let relative_history = Apply_expr.relative_history apply in
   let position = Apply_expr.position apply in
+  let region = Apply_expr.region apply in
   Apply_expr.create ~callee ~continuation exn_continuation ~args ~call_kind dbg
     ~inlined ~inlining_state ~probe_name:None ~position ~relative_history
+    ~region
   |> Expr.create_apply
 
 and subst_apply_cont env apply_cont =
@@ -898,7 +901,7 @@ let call_kinds env (call_kind1 : Call_kind.t) (call_kind2 : Call_kind.t) :
       ~subst2:(fun _ arity -> arity)
       env (code_id1, return_arity1) (code_id2, return_arity2)
     |> Comparison.map ~f:(fun (code_id, return_arity) ->
-           Call_kind.direct_function_call code_id ~return_arity Heap)
+           Call_kind.direct_function_call code_id ~return_arity Alloc_mode.heap)
   | ( Function
         { function_call =
             Indirect_known_arity
@@ -923,7 +926,7 @@ let call_kinds env (call_kind1 : Call_kind.t) (call_kind2 : Call_kind.t) :
     pairs ~f1:method_kinds ~f2:simple_exprs ~subst2:subst_simple env
       (kind1, obj1) (kind2, obj2)
     |> Comparison.map ~f:(fun (kind, obj) ->
-           Call_kind.method_call kind ~obj Heap)
+           Call_kind.method_call kind ~obj Alloc_mode.heap)
   | ( C_call
         { alloc = alloc1;
           param_arity = param_arity1;
@@ -988,6 +991,7 @@ let apply_exprs env apply1 apply2 : Expr.t Comparison.t =
             ~inlining_state:(Apply.inlining_state apply1)
             ~probe_name:None ~position:(Apply.position apply1)
             ~relative_history:(Apply_expr.relative_history apply1)
+            ~region:(Apply_expr.region apply1)
           |> Expr.create_apply
       }
 
@@ -1149,13 +1153,14 @@ and codes env (code1 : Code.t) (code2 : Code.t) =
            ~body1
            ~body2
            ~my_closure
+           ~my_region
            ~my_depth
          ->
         exprs env body1 body2
         |> Comparison.map ~f:(fun body1' ->
                Function_params_and_body.create ~return_continuation
-                 ~exn_continuation params ~body:body1' ~my_closure ~my_depth
-                 ~free_names_of_body:Unknown))
+                 ~exn_continuation params ~body:body1' ~my_closure ~my_region
+                 ~my_depth ~free_names_of_body:Unknown))
   in
   pairs ~f1:bodies
     ~f2:(options ~f:code_ids ~subst:subst_code_id)
@@ -1265,6 +1270,7 @@ and cont_handlers env handler1 handler2 =
 let flambda_units u1 u2 =
   let ret_cont = Continuation.create ~sort:Toplevel_return () in
   let exn_cont = Continuation.create () in
+  let toplevel_my_region = Variable.create "toplevel_my_region" in
   let mk_renaming u =
     let renaming = Renaming.empty in
     let renaming =
@@ -1277,6 +1283,11 @@ let flambda_units u1 u2 =
         (Flambda_unit.exn_continuation u)
         ~guaranteed_fresh:exn_cont
     in
+    let renaming =
+      Renaming.add_fresh_variable renaming
+        (Flambda_unit.toplevel_my_region u)
+        ~guaranteed_fresh:toplevel_my_region
+    in
     renaming
   in
   let env = Env.create () in
@@ -1287,4 +1298,4 @@ let flambda_units u1 u2 =
          let module_symbol = Flambda_unit.module_symbol u1 in
          Flambda_unit.create ~return_continuation:ret_cont
            ~exn_continuation:exn_cont ~body ~module_symbol
-           ~used_value_slots:Unknown)
+           ~used_value_slots:Unknown ~toplevel_my_region)