Get rid of unsafe_cleanup

lukstafi · lukstafi · commit 0a47d49bab3a · 2024-10-17T12:58:42.000+02:00
This is just a "dumb delete", might need a followup but nothing obvious. E.g. in the gccjit backend, `root_ctx` persists "forever".
diff --git a/CHANGES.md b/CHANGES.md
@@ -13,7 +13,8 @@
 - New files: split out `backend_types.ml` from `backends.ml`; moved `Tnode.task` to `task.ml`; renamed `backend_utils.ml` to `c_syntax.ml`.
 - Removed half-static verification of merge buffer nodes inside `device_to_device`.
 - Fixed #286: cross-stream-sharing incorporated into `Tnode.memory_mode`.
-- TODO: Moved the multicore backend from a `device = stream` model to a single device model.
+- Moved the multicore backend from a `device = stream` model to a single device model.
+- Got rid of `unsafe_cleanup`.
 - TODO: Built per-tensor-node stream-to-stream synchronization into device-to-device copying functions, removed obsolete blocking synchronizations.
 
 ### Fixed
diff --git a/arrayjit/lib/backend_types.ml b/arrayjit/lib/backend_types.ml
@@ -54,12 +54,11 @@ module type Backend_common = sig
   val name : string
 
   val initialize : Types.config -> unit
-  (** Initializes a backend before first use or (on some backends) after {!unsafe_cleanup}. Does
-      nothing if the backend is already initialized. *)
+  (** Initializes a backend before first use. Typically does nothing if the backend is already
+      initialized, but some backends can do some safe cleanups. *)
 
   val is_initialized : unit -> bool
-  (** Returns false if there was no previous {!initialize} call, or, on some backends, the most
-      recent call was followed by {!unsafe_cleanup}. If it returns false, one must call
+  (** Returns false if there was no previous {!initialize} call. If it returns false, one must call
       {!initialize} before using the backend. *)
 
   val init : init_info -> context
@@ -88,11 +87,6 @@ module type Backend_common = sig
       compile time and debugging convenience by generating fewer files -- ideally does not affect
       execution, but there can be backend-specific differences. Only array entries for which
       [occupancy] returns true are included. *)
-
-  val unsafe_cleanup : unit -> unit
-  (** Cleans up all work on a backend, releases resources. All previously retrieved values
-      (contexts, streams and devices) become invalid. The backend needs to be initialized again to
-      be used again. *)
 end
 
 module type No_device_backend = sig
@@ -236,7 +230,6 @@ module type Lowered_backend_common = sig
   val is_initialized : unit -> bool
   val init : init_info -> context
   val finalize : context -> unit
-  val unsafe_cleanup : unit -> unit
   val name : string
 end
 
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -320,10 +320,6 @@ struct
     Stdlib.Condition.broadcast stream.state.dev_wait_for_work;
     Domain.join stream.domain
 
-  let%track2_sexp unsafe_cleanup () =
-    latest_subordinal := 0;
-    Backend.unsafe_cleanup ()
-
   let get_device ~ordinal =
     if ordinal <> 0 then
       invalid_arg [%string "Multicore_backend.get_device %{ordinal#Int}: only device 0 exists"];
@@ -478,10 +474,6 @@ module Sync_backend (Backend : Backend_types.No_device_backend) : Backend_types.
   let suggested_num_streams _device = !sync_suggested_num_streams
   let next_stream_id = ref 0
 
-  let unsafe_cleanup () =
-    next_stream_id := 0;
-    Backend.unsafe_cleanup ()
-
   let get_device ~ordinal =
     if ordinal <> 0 then invalid_arg "Sync_backend backends only have device number 0";
     CPU
@@ -807,7 +799,6 @@ let reinitialize (module Backend : Backend_types.Backend) config =
   if not @@ Backend.is_initialized () then Backend.initialize config
   else (
     Core.Gc.full_major ();
-    Backend.unsafe_cleanup ();
     Backend.initialize config)
 
 (** Reinitializes and returns a backend corresponding to [backend_name], or if omitted, selected via
diff --git a/arrayjit/lib/cc_backend.ml b/arrayjit/lib/cc_backend.ml
@@ -40,7 +40,6 @@ let to_buffer tn ~dst ~src =
 
 let host_to_buffer src ~dst = Ndarray.map2 { f2 = Ndarray.A.blit } src dst
 let buffer_to_host dst ~src = Ndarray.map2 { f2 = Ndarray.A.blit } src dst
-let unsafe_cleanup () = Stdlib.Gc.compact ()
 
 let is_initialized, initialize =
   let initialized = ref false in
diff --git a/arrayjit/lib/cuda_backend.cudajit.ml b/arrayjit/lib/cuda_backend.cudajit.ml
@@ -235,18 +235,6 @@ let init stream =
   Stdlib.Gc.finalise finalize ctx;
   ctx
 
-let unsafe_cleanup () =
-  let len = Core.Weak.length !devices in
-  (* NOTE: releasing the context should free its resources, there's no need to finalize the
-     remaining contexts, and [finalize], [finalize_device] will not do anything for a [released]
-     device. *)
-  for i = 0 to len - 1 do
-    Option.iter (Core.Weak.get !devices i) ~f:(fun device ->
-        if Atomic.compare_and_set device.released false true then cleanup_device device)
-  done;
-  Core.Weak.fill !devices 0 len None;
-  Stdlib.Gc.compact ()
-
 let%diagn2_l_sexp from_host (ctx : context) tn =
   match (tn, Map.find ctx.ctx_arrays tn) with
   | { Tn.array = (lazy (Some hosted)); _ }, Some dst ->
diff --git a/arrayjit/lib/cuda_backend.missing.ml b/arrayjit/lib/cuda_backend.missing.ml
@@ -52,7 +52,6 @@ let link_batch (Unimplemented_ctx : context) (code_batch : code_batch) =
   in
   ((Unimplemented_ctx : context), lowered_bindings, task)
 
-let unsafe_cleanup () = ()
 let from_host _context _tn = false
 let to_host _context _tn = false
 let device_to_device _tn ~into_merge_buffer:_ ~dst:_ ~src:_ = false
diff --git a/arrayjit/lib/gcc_backend.gccjit.ml b/arrayjit/lib/gcc_backend.gccjit.ml
@@ -55,12 +55,6 @@ let to_buffer tn ~dst ~src =
 let host_to_buffer src ~dst = Ndarray.map2 { f2 = Ndarray.A.blit } src dst
 let buffer_to_host dst ~src = Ndarray.map2 { f2 = Ndarray.A.blit } src dst
 
-let unsafe_cleanup () =
-  let open Gccjit in
-  Option.iter ~f:Context.release !root_ctx;
-  root_ctx := None;
-  Stdlib.Gc.compact ()
-
 let is_initialized () = Option.is_some !root_ctx
 
 let initialize () =
diff --git a/arrayjit/lib/gcc_backend.missing.ml b/arrayjit/lib/gcc_backend.missing.ml
@@ -47,4 +47,3 @@ let initialize () = ()
 let is_initialized () = true
 let init _ = Unimplemented_ctx
 let finalize Unimplemented_ctx = ()
-let unsafe_cleanup () = ()
diff --git a/bin/compilation_speed.ml b/bin/compilation_speed.ml
@@ -76,8 +76,6 @@ let benchmark_overhead backend () =
   in
   PrintBox_text.output Stdio.stdout plot_box;
   Stdio.print_endline "\n";
-  let module Backend = (val backend) in
-  Backend.unsafe_cleanup ();
   result
 
 let benchmarks =
diff --git a/bin/einsum_trivia.ml b/bin/einsum_trivia.ml
@@ -39,8 +39,7 @@ let _suspended () =
   let%op ho2 = hey2 ++ "ab|cd->ef => cf|ae->db" in
   Utils.capture_stdout_logs @@ fun () ->
   Train.forward_and_forget backend ctx ho2;
-  Tensor.print ~force:true ~with_code:false ~with_grad:false `Default @@ ho2;
-  Backend.unsafe_cleanup ()
+  Tensor.print ~force:true ~with_code:false ~with_grad:false `Default @@ ho2
 
 let () =
   Utils.set_log_level 2;
diff --git a/bin/micrograd_demo.ml b/bin/micrograd_demo.ml
@@ -180,8 +180,7 @@ let experiment seed ~no_batch_shape_inference ~use_builtin_weight_decay () =
     plot ~size:(120, 30) ~x_label:"step" ~y_label:"learning rate"
       [ Line_plot { points = Array.of_list_rev !learning_rates; pixel = "-" } ]
   in
-  PrintBox_text.output Stdio.stdout plot_lr;
-  Backend.unsafe_cleanup ()
+  PrintBox_text.output Stdio.stdout plot_lr
 
 let () = experiment 4 ~no_batch_shape_inference:true ~use_builtin_weight_decay:true ()
 let () = experiment 4 ~no_batch_shape_inference:false ~use_builtin_weight_decay:false ()
diff --git a/bin/moons_benchmark.ml b/bin/moons_benchmark.ml
@@ -178,8 +178,6 @@ let classify_moons ~seed ~on_device ~inlining_cutoff ~num_streams ~batch_size ~b
       }
   in
   Stdio.printf "\n\n%!";
-  let module Backend = (val backend) in
-  Backend.unsafe_cleanup ();
   result
 
 let _suspend () =
diff --git a/bin/moons_demo.ml b/bin/moons_demo.ml
@@ -136,6 +136,6 @@ let demo () =
     PrintBox_text.output Stdio.stdout plot_moons;
     Stdio.print_endline "\n"
   in
-  Backend.unsafe_cleanup ()
+  ()
 
 let () = demo ()
diff --git a/bin/moons_demo_parallel.ml b/bin/moons_demo_parallel.ml
@@ -131,9 +131,7 @@ let experiment ~seed ~backend_name ~config () =
     plot ~size:(120, 30) ~x_label:"step" ~y_label:"learning rate"
       [ Line_plot { points = Array.of_list_rev learning_rates; pixel = "-" } ]
   in
-  PrintBox_text.output Stdio.stdout plot_lr;
-  let module Backend = (val backend) in
-  Backend.unsafe_cleanup ()
+  PrintBox_text.output Stdio.stdout plot_lr
 
 let () = experiment ~seed:1 ~backend_name:"cc" ~config:Only_devices_parallel ()
 let _suspended () = experiment ~seed:1 ~backend_name:"cc" ~config:Only_devices_parallel ()
diff --git a/bin/zero2hero_1of7.ml b/bin/zero2hero_1of7.ml
@@ -145,9 +145,6 @@ let _suspended () =
     in
     PrintBox_text.output Stdio.stdout plot_box
   in
-  (* Test successive cleanups. *)
-  Backend.unsafe_cleanup ();
-  Backend.unsafe_cleanup ();
   ()
 
 let _suspended () =
