Attempt to shrink the heap before calling the assembler (#429)

* Attempt to shrink the heap before calling the assembler

We have a fair amount of global state that's never collected, which is causing
trouble in large parallel builds, particularly when large assembly files are
generated, causing `as` to need a lot of memory itself. This patch simply clears
out some data once it's no longer needed and, when memory use is high, calls
`Gc.compact` to try and release some memory back to the OS. This has been
observed to reduce the heap size by some 640MB (for an unusually large generated
.ml file).

This PR still leaves some relatively low-hanging fruit: after to_cmm, we can
marshal the .cmx eagerly and clear out the exported info. However, it's not
obvious how to handle this cleanly, and anyway one can use `-Oclassic` to avoid
producing a large .cmx to begin with.

Author: lukemaurer

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 ()