- Stream-to-stream synchronization functionality, with lazy per-tensor-node synchronization.

- TODO: Automatic blocking on access of a host array when a scheduled `to_host` transfer has not finished.

- Built per-tensor-node stream-to-stream synchronization into copying functions.

- Re-introduced whole-device blocking synchronization, which now is just a slight optimization as it also cleans up event book-keeping.

streams = Utils.weak_create ();

let make_stream device runner =

Utils.register_new device.streams ~grow_by:8 (fun stream_id ->

{

device;

runner;

merge_buffer = ref None;

stream_id;

allocated_buffer = None;

updating_for = Hashtbl.create (module Tnode);

updating_for_merge_buffer = None;

reader_streams = Hashtbl.create (module Tnode);

})

mutable streams : ('buffer_ptr, 'dev, 'runner, 'event) stream_ref Utils.weak_dynarray;

mutable streams : ('buffer_ptr, 'dev, 'runner, 'event) stream_ref Utils.weak_dynarray; (** . *)

stream_id : int; (** An ID unique within the device for the lifetime of the stream. *)

val make_stream : device -> runner -> stream

type device

val sync_device : device -> unit

(** Synchronizes all the streams on a device, and cleans up (removes) all associated events. *)

include

With_buffer_retrieval_and_syncing

with type device := device

and type context := context

and type event := event

let[@landmark] wait_for_all ctx streams tn =

let[@landmark] wait_for_ready ~dst ~src tn =

let[@landmark] update_writer_event ?e ?from s tn =

(s, e) :: Option.value ~default:[] l)

else Hashtbl.remove s.device.shared_writer_streams tn;

let%track2_l_sexp[@landmark] from_host (ctx : Backend.context) tn =

let%track2_l_sexp[@landmark] to_host (ctx : Backend.context) (tn : Tn.t) =

let%diagn2_l_sexp[@landmark] device_to_device (tn : Tn.t) ~into_merge_buffer

~(dst : Backend.context) ~(src : Backend.context) =

let[@landmark] merge_task () = Task.run task in

merge_task ();

let%track2_l_sexp sync_routine r =

let[@landmark] pre () =

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

let data = List.filter data ~f:(fun (_, e) -> not (Backend.is_done e)) in

Hashtbl.set s.device.shared_writer_streams ~key:tn ~data;

List.iter data ~f:(fun (work_stream, e) ->

if not (equal_stream work_stream s) then Backend.will_wait_for r.context e))

else Hashtbl.remove s.device.shared_writer_streams tn)

let[@landmark] post () =

let[@landmark] sync_device device =

Utils.weak_iter device.streams ~f:Backend.await;

Hashtbl.clear device.host_writing_streams;

Hashtbl.clear device.host_reading_streams;

Hashtbl.clear device.shared_writer_streams;

Utils.weak_iter device.streams ~f:(fun s ->

Hashtbl.clear s.reader_streams;

s.updating_for_merge_buffer <- None;

Hashtbl.clear s.updating_for)

let[@landmark] from_host ~dst_ptr ~dst hosted =

let[@landmark] to_host ~src_ptr ~src hosted =

let[@landmark] device_to_device tn ~into_merge_buffer ~dst_ptr ~dst ~src_ptr ~src =

let[@landmark] finalize (type buffer_ptr dev runner event)

(fun ~message:_message ~status:_status ->

_message;

if not @@ Cu.is_success _status then [%log (_status : Cu.result)]]])

make_stream device cu_stream

let[@landmark] await stream : unit =

let[@landmark] from_host ~dst_ptr ~dst hosted =

let[@landmark] to_host ~src_ptr ~src hosted =

let[@landmark] device_to_device tn ~into_merge_buffer ~dst_ptr ~dst ~src_ptr ~src =

(* FIXME: every now and then the compilation crashes because the options are garbled. *)

(* Stdio.printf "PTX options %s\n%!" @@ String.concat ~sep:", " options; *)

landmarks

landmarks-ppx