Turn off automatic host transfers on demand

lukstafi · lukstafi · commit a58eabb433e8 · 2024-12-31T20:01:13.000+01:00
diff --git a/CHANGES.md b/CHANGES.md
@@ -1,4 +1,4 @@
-## [0.5.1] -- current
+## [0.5.1] -- 2024-12-31
 
 ## Added
 
@@ -7,7 +7,7 @@
 
 ## Fixed
 
-- Added `#` as alternative to `~~` for comment lines in `ocannl_config` files, and a bug in parsing.
+- Added `#` as alternative to `~~` for comment lines in `ocannl_config` files, and fixed a bug in their parsing.
 
 ## [0.5.0] -- 2024-12-18
 
diff --git a/arrayjit/lib/backends.ml b/arrayjit/lib/backends.ml
@@ -144,7 +144,9 @@ module Add_buffer_retrieval_and_syncing (Backend : No_buffer_retrieval_or_syncin
     let hosted_inputs = Set.filter r.inputs ~f:(fun tn -> Tn.is_hosted_force tn 47) in
     let pre () =
       assert (Domain.is_main_domain ());
-      Set.iter hosted_inputs ~f:(fun tn -> if tn.host_modified then assert (from_host r.context tn));
+      if Utils.settings.automatic_host_transfers then
+        Set.iter hosted_inputs ~f:(fun tn ->
+            if tn.host_modified then assert (from_host r.context tn));
       Set.iter r.inputs ~f:(fun tn ->
           if Tn.potentially_cross_stream tn then
             Option.iter (Hashtbl.find s.device.shared_writer_streams tn) ~f:(fun data ->
diff --git a/arrayjit/lib/tnode.ml b/arrayjit/lib/tnode.ml
@@ -589,7 +589,7 @@ let do_read tn =
   Option.iter
     ~f:(fun p ->
       p.sync ();
-      p.transfer ())
+      if Utils.settings.automatic_host_transfers then p.transfer ())
     tn.prepare_read;
   tn.prepare_read <- None
 
diff --git a/arrayjit/lib/utils.ml b/arrayjit/lib/utils.ml
@@ -42,6 +42,13 @@ type settings = {
   mutable check_half_prec_constants_cutoff : float option;
       (** If given, generic code optimization should fail if a half precision FP16 constant exceeds
           the cutoff. *)
+  mutable automatic_host_transfers : bool;
+      (** If true, [from_host] and [to_host] happen automatically in specific situations.
+          - When a host array is about to be read, we transfer to host from the context that most
+            recently updated the node.
+          - When a routine is about to be run, we transfer the routine's inputs from host to the
+            routine's context if the host array was not yet transfered since its creation or most
+            recent modification. *)
 }
 [@@deriving sexp]
 
@@ -53,6 +60,7 @@ let settings =
     fixed_state_for_init = None;
     print_decimals_precision = 2;
     check_half_prec_constants_cutoff = Some (2. **. 14.);
+    automatic_host_transfers = true;
   }
 
 let accessed_global_args = Hash_set.create (module String)
@@ -364,7 +372,9 @@ let restore_settings () =
     Int.of_string @@ get_global_arg ~arg_name:"print_decimals_precision" ~default:"2";
   settings.check_half_prec_constants_cutoff <-
     Float.of_string_opt
-    @@ get_global_arg ~arg_name:"check_half_prec_constants_cutoff" ~default:"16384.0"
+    @@ get_global_arg ~arg_name:"check_half_prec_constants_cutoff" ~default:"16384.0";
+  settings.automatic_host_transfers <-
+    Bool.of_string @@ get_global_arg ~arg_name:"automatic_host_transfers" ~default:"true"
 
 let () = restore_settings ()
 let with_runtime_debug () = settings.output_debug_files_in_build_directory && settings.log_level > 1
diff --git a/ocannl_config.example b/ocannl_config.example
@@ -35,6 +35,14 @@ print_decimals_precision=2
 # Complains if a half-precision tensor node is a constant with absolute value exceeding this.
 check_half_prec_constants_cutoff=16384.0
 
+# If true, [from_host] and [to_host] happen automatically in specific situations:
+# - When a host array is about to be read, we transfer to host from the context that most
+#   recently updated the node.
+# - When a routine is about to be run, we transfer the routine's inputs from host to the
+#   routine's context if the host array was not yet transfered since its creation or most
+#   recent modification.
+automatic_host_transfers=true
+
 # Other configurations:
 
 # If true, stdout capturing is disabled, so some logs meant for the ppx_minidebug log files
diff --git a/todo.md b/todo.md
@@ -1,5 +1,5 @@
 # This file is for tasks with a smaller granularity than issues, typically immediate tasks.
-(A) Ensure that reading from host on CPU performs required synchronization
+(A) Ensure that reading from host on CPU performs required synchronization {cm:2024-12-31}
 
 Update `anatomy_of_a_backend.md`
 Update introductory slides {cm:2024-12-17}
