OCANNL expects backends to implement FIFO queue scheduling, and an event mechanism for synchronizing between streams (and ideally devices), matching the CUDA specification. On top of events, OCANNL implements per-tensor-node synchronization. 1/3rd of the `device` fields have to do with synchronization:

Besides routines, calling `from_host`, `to_host`, `device_to_device` from a backend puts the corresponding tasks on the device's queue. Both invoking a routine and calling these copying functions will perform the necessary event creations and synchronizations to ensure that when scheduling writing into an array precedes scheduling reading from it, the actual writing also precedes the actual reading.

{

dev;

ordinal;

latest_stream_id = -1;

released = Atomic.make false;

cross_stream_candidates = Hashtbl.create (module Tnode);

device;

runner;

merge_buffer = ref None;

stream_id;

allocated_buffer = None;

updating_for = Hashtbl.create (module Tnode);

(** Like {!Backend.from_host}, but without synchronization and buffer retrieval. *)

val to_host : src_ptr:buffer_ptr -> src:context -> Ndarray.t -> unit

val device_to_device :

Tnode.t ->

src_ptr:buffer_ptr ->

src:context ->

unit

inputs : Set.M(Tnode).t;

(** The materialized read-only and read-before-write (within the routine) non-constant nodes.

outputs : Set.M(Tnode).t; (** All the materialized nodes written-to by the routine. *)

}

[@@deriving sexp_of]

type ('buffer_ptr, 'dev, 'runner, 'event) device_ref = {

dev : 'dev;

ordinal : int;

mutable latest_stream_id : int;

released : Utils.atomic_bool;

cross_stream_candidates : 'buffer_ptr Hashtbl.M(Tnode).t;

device : ('buffer_ptr, 'dev, 'runner, 'event) device_ref;

runner : 'runner;

merge_buffer : 'buffer_ptr buffer option ref;

stream_id : int;

mutable allocated_buffer : 'buffer_ptr buffer option;

updating_for : 'event Hashtbl.M(Tnode).t;

('buffer_ptr, 'dev, 'runner, 'event) device_ref = {

dev : 'dev;

ordinal : int;

mutable latest_stream_id : int;

released : Utils.atomic_bool;

cross_stream_candidates : 'buffer_ptr Hashtbl.M(Tnode).t;

runner : 'runner;

merge_buffer : 'buffer_ptr buffer option ref;

(** Depending on backend implementations, either the currently used merge buffer, or the one

stream_id : int; (** An ID unique within the device. *)

mutable allocated_buffer : 'buffer_ptr buffer option;

updating_for : 'event Hashtbl.M(Tnode).t;

mutable updating_for_merge_buffer : (Tnode.t * 'event) option;

reader_streams :

(('buffer_ptr, 'dev, 'runner, 'event) stream_ref * 'event) list Hashtbl.M(Tnode).t;

(** The streams, other than this stream, that most recently have been reading from a node in

let check_merge_buffer stream ~code_node =

let name = function Some tn -> Tnode.debug_name tn | None -> "none" in

| _, None -> ()

| _ ->

raise

@@ Utils.User_error

module Add_buffer_retrieval_and_syncing (Backend : No_buffer_retrieval_or_syncing) = struct

let wait_for_all ctx streams tn =

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

Hashtbl.update s.updating_for tn ~f:(fun _ -> e)

| Merge_buffer tn ->

s.updating_for_merge_buffer <- Some (tn, e)

let%diagn2_l_sexp from_host (ctx : Backend.context) tn =

update_writer_event ~from:(`Src src.stream) dst.stream @@ Node tn;

[%log "copying", Tn.debug_name tn, "from", name_of src, "to", name_of dst];

true)

| Copy | Streaming ->

Backend.(

device_to_device tn ~into_merge_buffer ~dst_ptr:None ~dst ~src_ptr:s_arr ~src);

update_writer_event ~from:(`Src src.stream) dst.stream @@ Merge_buffer tn;

true)

let%track3_l_sexp sync_routine r =

let s = r.context.stream in

let pre () =

in

let post () =

let e = Backend.all_work s in

let device_to_device tn ~into_merge_buffer ~dst_ptr ~dst ~src_ptr ~src =

let dev = dst.stream in

let size_in_bytes = Tnode.size_in_bytes tn in

let work =

(* TODO: log the operation if [Utils.settings.with_log_level > 1]. *)

match (into_merge_buffer, dst_ptr) with

"device_to_device " ^ Tnode.debug_name tn ^ " dst " ^ get_name dev ^ " src "

^ get_name src.stream

in

schedule_task dev (Task.Task { context_lifetime = (src, dst); description; work })

end

let free, total = Cudajit.Device.get_free_and_total_mem () in

total - free

if

buffer.size_in_bytes < size_in_bytes)

then (

let%track3_sexp cleanup_device (device : device) =

Cu.Context.set_current device.dev.primary_context;

Cu.Context.synchronize ();

Option.iter !Utils.advance_captured_logs ~f:(fun callback -> callback ());

(* Note: this is not necessary as releasing the primary context by GC will reset the context. *)

Hashtbl.iter device.cross_stream_candidates ~f:(fun buffer_ptr ->

Cu.Deviceptr.mem_free buffer_ptr)

if Utils.debug_log_from_routines () && not (Hash_set.mem initialized_devices ordinal) then

Option.iter Utils.settings.cuda_printf_fifo_size ~f:Cu.Context.(set_limit PRINTF_FIFO_SIZE);

Hash_set.add initialized_devices ordinal;

let result = make_device dev ~ordinal in

Stdlib.Gc.finalise finalize_device result;

Stdlib.Weak.set !devices ordinal (Some result);

memcpy ~dst_ptr

| Streaming, _ ->

assert same_device;

dst.stream.merge_buffer := Some { ptr = src_ptr; size_in_bytes }

| Copy, _ ->

set_ctx @@ ctx_of dst;

let size_in_bytes = Tn.size_in_bytes tn in

type code = {

traced_store : Low_level.traced_store;

~f:(visit ~is_assigned:(traced.zeroed_out || Hash_set.mem traced.assignments at_pos))

| Local_scope { body; _ } -> loop_proc env body

| Get_local _ -> ()

| Get_global _ -> ()

| Embed_index _ -> ()

| Binop (Arg1, llv1, _llv2) -> loop llv1

else outputs

in

(inputs, outputs)),

let%diagn2_sexp optimize_proc static_indices llc =

let traced_store = Hashtbl.create (module Tnode) in

t ->

optimized

assert (Backend.to_host sgd_update.context learning_rate.value);

(* scalar_loss is not in the sgd_update context. *)

assert (Backend.to_host grad_updates.(0).context scalar_loss.value);

let batch_loss = scalar_loss.@[0] in

epoch_loss := !epoch_loss +. batch_loss;

batch_losses := batch_loss :: !batch_losses;

Option.iter per_epoch_callback ~f:(fun f ->

f ~at_step:!step_ref ~at_epoch:epoch ~learning_rate:learning_rate.@[0]

~epoch_loss:!epoch_loss);

Array.iter streams ~f:(fun s ->

Stdlib.Format.printf "Stream %d debug %s:@ %a\n%!" s.stream_id pos Sexp.pp_hum

@@ Backend.get_debug_info s)

in

done;

let%op model_result = model "infer" in

let infer_fwd =