Attempt to shrink the heap before calling the assembler

backend/amd64/emit.mlp

@@ -1204,13 +1204,17 @@ let data l =
 
 (* Beginning / end of an assembly file *)
 
-let begin_assembly() =
+let reset_all () =
   X86_proc.reset_asm_code ();
+  Emitaux.reset ();
   reset_debug_info();                   (* PR#5603 *)
   reset_imp_table();
   reset_probes ();
   float_constants := [];
-  all_functions := [];
+  all_functions := []
+
+let begin_assembly() =
+  reset_all ();
   if system = S_win64 then begin
     D.extrn "caml_call_gc" NEAR;
     D.extrn "caml_c_call" NEAR;
@@ -1592,4 +1596,5 @@ let end_assembly() =
     else
       None
   in
-  X86_proc.generate_code asm
+  X86_proc.generate_code asm;
+  reset_all ()

backend/asmgen.ml

@@ -285,6 +285,11 @@ let compile_unit ~output_prefix ~asm_filename ~keep_asm ~obj_filename gen =
          ~exceptionally:(fun () ->
              if create_asm && not keep_asm then remove_file asm_filename);
        if should_emit () then begin
+         Emitaux.reduce_heap_size ~reset:(fun () ->
+            reset ();
+            Typemod.reset ();
+            Emitaux.reset ();
+            Reg.reset ());
          let assemble_result =
            Profile.record "assemble"
              (Proc.assemble_file asm_filename) obj_filename
@@ -344,9 +349,9 @@ let compile_implementation ?toplevel ~backend ~filename ~prefixname ~middle_end
       in
       end_gen_implementation ?toplevel ~ppf_dump clambda_with_constants)
 
-let compile_implementation_flambda2 ?toplevel ~filename ~prefixname
-    ~size:module_block_size_in_words ~module_ident ~module_initializer
-    ~flambda2 ~ppf_dump ~required_globals () =
+let compile_implementation_flambda2 ?toplevel ?(keep_symbol_tables=true)
+    ~filename ~prefixname ~size:module_block_size_in_words ~module_ident
+    ~module_initializer ~flambda2 ~ppf_dump ~required_globals () =
   compile_unit ~output_prefix:prefixname
     ~asm_filename:(asm_filename prefixname) ~keep_asm:!keep_asm_file
     ~obj_filename:(prefixname ^ ext_obj)
@@ -355,6 +360,7 @@ let compile_implementation_flambda2 ?toplevel ~filename ~prefixname
       let cmm_phrases =
         flambda2 ~ppf_dump ~prefixname ~filename ~module_ident
           ~module_block_size_in_words ~module_initializer
+          ~keep_symbol_tables
       in
       end_gen_implementation0 ?toplevel ~ppf_dump (fun () -> cmm_phrases))
 

backend/asmgen.mli

@@ -40,6 +40,7 @@ val compile_implementation
     Cmmgen pass.  Instead it emits Cmm directly. *)
 val compile_implementation_flambda2
    : ?toplevel:(string -> bool)
+  -> ?keep_symbol_tables:bool
   -> filename:string
   -> prefixname:string
   -> size:int
@@ -52,6 +53,7 @@ val compile_implementation_flambda2
     module_ident:Ident.t ->
     module_block_size_in_words:int ->
     module_initializer:Lambda.lambda ->
+    keep_symbol_tables:bool ->
     Cmm.phrase list)
   -> ppf_dump:Format.formatter
   -> required_globals:Ident.Set.t

backend/asmlink.ml

@@ -396,20 +396,7 @@ let link ~ppf_dump objfiles output_name =
       ~asm_filename:startup ~keep_asm:!Clflags.keep_startup_file
       ~obj_filename:startup_obj
       (fun () -> make_startup_file ~ppf_dump units_tolink);
-    (* Clear all state and compact before calling the linker, because the linker
-       can use a lot of memory, and this reduces the peak memory usage by freeing
-       most of the memory from this process before the linker starts using memory.
-
-       On a link where this frees up around 1.1GB of memory this takes around 0.6s. We
-       only take this time on large links where the number of parallel linking jobs is
-       likely to be constrained by total system memory. *)
-    let _minor, _promoted, major_words = Gc.counters () in
-    (* Uses [major_words] because it doesn't require a heap traversal to compute and 
-       for this workload a majority of major words are live at this point. *)
-    if major_words > 500_000_000.0 /. 8.0 then 
-      Profile.record_call "asmlink_compact" (fun () ->
-        reset ();
-        Gc.compact ());
+    Emitaux.reduce_heap_size ~reset:(fun () -> reset ());
     Misc.try_finally
       (fun () ->
          call_linker (List.map object_file_name objfiles)

backend/asmpackager.ml

@@ -114,6 +114,7 @@ let make_package_object ~ppf_dump members targetobj targetname coercion
         ~flambda2
         ~ppf_dump
         ~required_globals:required_globals
+        ~keep_symbol_tables:false
         ()
     end else begin
       let program, middle_end =

backend/asmpackager.mli

@@ -29,6 +29,7 @@ val package_files
     module_ident:Ident.t ->
     module_block_size_in_words:int ->
     module_initializer:Lambda.lambda ->
+    keep_symbol_tables:bool ->
     Cmm.phrase list)
   -> unit
 

backend/emitaux.ml

@@ -375,3 +375,13 @@ let reset () =
 
 let binary_backend_available = ref false
 let create_asm_file = ref true
+
+let reduce_heap_size ~reset =
+  let _minor, _promoted, major_words = Gc.counters () in
+  (* Uses [major_words] because it doesn't require a heap traversal to compute and 
+     for this workload a majority of major words are live at this point. *)
+  if major_words > 500_000_000.0 /. 8.0 then begin
+    Profile.record_call "compact" (fun () ->
+      reset ();
+      Gc.compact ())
+  end

backend/emitaux.mli

@@ -81,3 +81,18 @@ val binary_backend_available: bool ref
 
 val create_asm_file: bool ref
     (** Are we actually generating the textual assembly file? *)
+
+(** Clear global state and compact the heap, so that an external program
+    (such as the assembler or linker) may have more memory available to it.
+
+    When this frees up around 1.1GB of memory, it takes around 0.6s. We only
+    take this time when the job is large enough that we're worried that we'll
+    either run out of memory or constrain the number of parallel jobs. We
+    heuristically measure how big the job is by how much heap we're using
+    ourselves.
+
+    The [reset] parameter will be called before [Gc.compact] if we go ahead
+    with the compaction. It should clear as much as possible from the global
+    state, since the fewer live words there are after GC, the smaller the new
+    heap can be. *)
+val reduce_heap_size : reset:(unit -> unit) -> unit

driver/optcompile.ml

@@ -51,7 +51,7 @@ let flambda_and_flambda2 i typed ~compile_implementation =
              ~required_globals;
            Compilenv.save_unit_info (cmx i)))
 
-let flambda2_ i ~flambda2 typed =
+let flambda2_ i ~flambda2 ~keep_symbol_tables typed =
   flambda_and_flambda2 i typed
     ~compile_implementation:(fun ~module_ident ~main_module_block_size ~code
           ~required_globals ->
@@ -64,6 +64,7 @@ let flambda2_ i ~flambda2 typed =
         ~flambda2
         ~ppf_dump:i.ppf_dump
         ~required_globals
+        ~keep_symbol_tables
         ())
 
 let flambda i backend typed =
@@ -112,13 +113,14 @@ let emit i =
   Compilenv.reset ?packname:!Clflags.for_package i.module_name;
   Asmgen.compile_implementation_linear i.output_prefix ~progname:i.source_file
 
-let implementation ~backend ~flambda2 ~start_from ~source_file ~output_prefix =
+let implementation ~backend ~flambda2 ~start_from ~source_file ~output_prefix
+    ~keep_symbol_tables =
   let backend info typed =
     Compilenv.reset ?packname:!Clflags.for_package info.module_name;
     if Config.flambda
     then flambda info backend typed
     else if Config.flambda2
-    then flambda2_ info ~flambda2 typed
+    then flambda2_ info ~flambda2 ~keep_symbol_tables typed
     else clambda info backend typed
   in
   with_info ~source_file ~output_prefix ~dump_ext:"cmx" @@ fun info ->

driver/optcompile.mli

@@ -26,9 +26,11 @@ val implementation
     module_ident:Ident.t ->
     module_block_size_in_words:int ->
     module_initializer:Lambda.lambda ->
+    keep_symbol_tables:bool ->
     Cmm.phrase list)
   -> start_from:Clflags.Compiler_pass.t
-  -> source_file:string -> output_prefix:string -> unit
+  -> source_file:string -> output_prefix:string -> keep_symbol_tables:bool
+  -> unit
 
 (** {2 Internal functions} **)
 

driver/optmaindriver.mli

@@ -28,5 +28,6 @@ val main
     module_ident:Ident.t ->
     module_block_size_in_words:int ->
     module_initializer:Lambda.lambda ->
+    keep_symbol_tables:bool ->
     Cmm.phrase list)
   -> int

middle_end/compilenv.ml

@@ -36,6 +36,10 @@ let global_infos_table =
 let export_infos_table =
   (Hashtbl.create 10 : (string, Export_info.t) Hashtbl.t)
 
+let reset_info_tables () =
+  Hashtbl.reset global_infos_table;
+  Hashtbl.reset export_infos_table
+
 let imported_sets_of_closures_table =
   (Set_of_closures_id.Tbl.create 10
    : Simple_value_approx.function_declarations option

middle_end/compilenv.mli

@@ -33,6 +33,8 @@ val reset: ?packname:string -> string -> unit
         (* Reset the environment and record the name of the unit being
            compiled (arg).  Optional argument is [-for-pack] prefix. *)
 
+val reset_info_tables: unit -> unit
+
 val unit_id_from_name: string -> Ident.t
         (* flambda-only *)
 

middle_end/flambda2/flambda2.ml

@@ -144,7 +144,7 @@ let output_flexpect ~ml_filename ~raw_flambda:old_unit new_unit =
         Format.pp_print_flush ppf ())
 
 let lambda_to_cmm ~ppf_dump:ppf ~prefixname ~filename ~module_ident
-    ~module_block_size_in_words ~module_initializer =
+    ~module_block_size_in_words ~module_initializer ~keep_symbol_tables =
   Misc.Color.setup (Flambda_features.colour ());
   (* When the float array optimisation is enabled, the length of an array needs
      to be computed differently according to the array kind, in the case where
@@ -203,7 +203,17 @@ let lambda_to_cmm ~ppf_dump:ppf ~prefixname ~filename ~module_ident
               Flambda2_identifiers.Code_id.print (Code.code_id code)
               Cost_metrics.print size)
     end;
-    Flambda2_to_cmm.To_cmm.unit ~make_symbol:Compilenv.make_symbol flambda cmx
-      ~all_code
+    let cmm =
+      Flambda2_to_cmm.To_cmm.unit ~make_symbol:Compilenv.make_symbol flambda cmx
+        ~all_code
+    in
+    if not keep_symbol_tables
+    then begin
+      Compilenv.reset_info_tables ();
+      Flambda2_identifiers.Code_id.reset ();
+      Flambda2_identifiers.Continuation.reset ();
+      Flambda2_identifiers.Reg_width_things.reset ()
+    end;
+    cmm
   in
   Profile.record_call "flambda2" run

middle_end/flambda2/flambda2.mli

@@ -26,6 +26,7 @@ val lambda_to_cmm :
   module_ident:Ident.t ->
   module_block_size_in_words:int ->
   module_initializer:Lambda.lambda ->
+  keep_symbol_tables:bool ->
   Cmm.phrase list
 
 val symbol_for_global :

middle_end/flambda2/identifiers/code_id.ml

@@ -67,6 +67,8 @@ let grand_table_of_code_ids = ref (Table.create ())
 
 let initialise () = grand_table_of_code_ids := Table.create ()
 
+let reset () = initialise ()
+
 let find_data t = Table.find !grand_table_of_code_ids t
 
 let get_compilation_unit t = (find_data t).compilation_unit

middle_end/flambda2/identifiers/code_id.mli

@@ -24,6 +24,8 @@ module Lmap : Lmap.S with type key = t
 
 val initialise : unit -> unit
 
+val reset : unit -> unit
+
 val create : name:string -> Compilation_unit.t -> t
 
 val get_compilation_unit : t -> Compilation_unit.t

middle_end/flambda2/identifiers/continuation.ml

@@ -137,6 +137,8 @@ let grand_table_of_continuations = ref (Table.create ())
 
 let initialise () = grand_table_of_continuations := Table.create ()
 
+let reset () = initialise ()
+
 (* CR mshinwell: Document why this uses [next_raise_count]. Does it need to? It
    would be better if it didn't. *)
 let create ?sort ?name () : t =

middle_end/flambda2/identifiers/continuation.mli

@@ -58,3 +58,5 @@ module With_args : sig
 end
 
 val initialise : unit -> unit
+
+val reset : unit -> unit

middle_end/flambda2/identifiers/reg_width_things.ml

@@ -730,3 +730,5 @@ let initialise () =
   Variable.initialise ();
   Symbol.initialise ();
   Simple.initialise ()
+
+let reset () = initialise ()

middle_end/flambda2/identifiers/reg_width_things.mli

@@ -198,3 +198,5 @@ module Simple : sig
 end
 
 val initialise : unit -> unit
+
+val reset : unit -> unit

ocaml/driver/compenv.ml

@@ -621,11 +621,13 @@ let c_object_of_filename name =
   Filename.chop_suffix (Filename.basename name) ".c" ^ Config.ext_obj
 
 let process_action
-    (ppf, implementation, interface, ocaml_mod_ext, ocaml_lib_ext) action =
+    (ppf, implementation, interface, ocaml_mod_ext, ocaml_lib_ext) action
+    ~keep_symbol_tables =
   let impl ~start_from name =
     readenv ppf (Before_compile name);
     let opref = output_prefix name in
-    implementation ~start_from ~source_file:name ~output_prefix:opref;
+    implementation ~start_from ~source_file:name ~output_prefix:opref
+      ~keep_symbol_tables;
     objfiles := (opref ^ ocaml_mod_ext) :: !objfiles
   in
   match action with
@@ -709,7 +711,14 @@ let process_deferred_actions env =
       | ProcessOtherFile name -> Filename.check_suffix name ".cmxa"
       | _ -> false) !deferred_actions then
     fatal "Option -a cannot be used with .cmxa input files.";
-  List.iter (process_action env) (List.rev !deferred_actions);
+  let compiling_multiple_impls =
+    List.length (List.filter (function
+        | ProcessImplementation _ -> true
+        | _ -> false) !deferred_actions) > 1
+  in
+  let keep_symbol_tables = compiling_multiple_impls in
+  List.iter (process_action env ~keep_symbol_tables)
+    (List.rev !deferred_actions);
   output_name := final_output_name;
   stop_early :=
     !compile_only ||

ocaml/driver/compenv.mli

@@ -82,7 +82,8 @@ val intf : string -> unit
 val process_deferred_actions :
   Format.formatter *
   (start_from:Clflags.Compiler_pass.t ->
-   source_file:string -> output_prefix:string -> unit) *
+   source_file:string -> output_prefix:string ->
+   keep_symbol_tables:bool -> unit) *
   (* compile implementation *)
   (source_file:string -> output_prefix:string -> unit) *
   (* compile interface *)

ocaml/driver/compile.ml

@@ -57,7 +57,8 @@ let emit_bytecode i (bytecode, required_globals) =
          (Emitcode.to_file oc i.module_name cmofile ~required_globals);
     )
 
-let implementation ~start_from ~source_file ~output_prefix =
+let implementation ~start_from ~source_file ~output_prefix
+    ~keep_symbol_tables:_ =
   let backend info typed =
     let bytecode = to_bytecode info typed in
     emit_bytecode info bytecode

ocaml/driver/compile.mli

@@ -19,7 +19,7 @@ val interface:
   source_file:string -> output_prefix:string -> unit
 val implementation:
   start_from:Clflags.Compiler_pass.t ->
-  source_file:string -> output_prefix:string -> unit
+  source_file:string -> output_prefix:string -> keep_symbol_tables:bool -> unit
 
 (** {2 Internal functions} **)
 

ocaml/driver/optcompile.ml

@@ -87,7 +87,8 @@ let emit i =
   Compilenv.reset ?packname:!Clflags.for_package i.module_name;
   Asmgen.compile_implementation_linear i.output_prefix ~progname:i.source_file
 
-let implementation ~backend ~start_from ~source_file ~output_prefix =
+let implementation ~backend ~start_from ~source_file
+    ~output_prefix ~keep_symbol_tables:_ =
   let backend info typed =
     Compilenv.reset ?packname:!Clflags.for_package info.module_name;
     if Config.flambda

ocaml/driver/optcompile.mli

@@ -20,7 +20,8 @@ val interface: source_file:string -> output_prefix:string -> unit
 val implementation:
    backend:(module Backend_intf.S)
    -> start_from:Clflags.Compiler_pass.t
-   -> source_file:string -> output_prefix:string -> unit
+   -> source_file:string -> output_prefix:string -> keep_symbol_tables:bool
+   -> unit
 
 (** {2 Internal functions} **)
 

ocaml/typing/typemod.ml

@@ -2943,3 +2943,8 @@ let () =
       | _ ->
         None
     )
+
+let reset () =
+  Env.reset_cache ();
+  Envaux.reset_cache ();
+  Typetexp.reset_type_variables ()