- Fixed #286: cross-stream-sharing incorporated into `Tnode.memory_mode`.

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

val get_ctx_stream : context -> stream

val get_stream_device : stream -> device

val to_ordinal : device -> int

val get_name : stream -> string

val get_ctx_stream : context -> stream

val get_name : stream -> string

val to_ordinal : device -> int

("Merge buffer mismatch, on stream: " ^ name scheduled_node ^ ", expected by code: "

^ name code_node)

module Domain = Domain [@warning "-3"]

state : stream_state;

merge_buffer : (buffer_ptr * Tnode.t) option ref;

mutable allocated_buffer : (buffer_ptr * int) option;

domain : (unit Domain.t[@sexp.opaque]);

}

[@@deriving sexp_of]

Stdlib.Condition.wait d.host_wait_for_idle d.mut

done;

Mut.unlock d.mut;

(** TODO: Returns the event indicating if any currently running or scheduled computations on the

let all_work _stream = Not_implemented_yet

let%track3_l_sexp schedule_task stream task =

assert (Domain.is_main_domain ());

let d = stream.state in

if not d.keep_spinning then invalid_arg "Multicore_backend: stream not available";

if not @@ Queue.try_push d.queue task then (

await stream;

let global_run_no = ref 0

Int.incr global_run_no;

let state =

{

with e ->

state.stream_error <- Some e;

state.keep_spinning <- false;

(* TODO: we risk raising this error multiple times because await and schedule_task raise

raise e

in

{

state;

domain = Domain.spawn worker;

merge_buffer = ref None;

allocated_buffer = None;

type context = { stream : stream; ctx : Backend.context } [@@deriving sexp_of]

type nonrec routine = context routine [@@deriving sexp_of]

let initialize = Backend.initialize

let is_initialized = Backend.is_initialized

let compile = Backend.compile

let compile_batch = Backend.compile_batch

let link { ctx; stream } code =

let task = Backend.link ~merge_buffer:stream.merge_buffer ctx code in

{

task with

context = { ctx = task.context; stream };

}

let link_batch { ctx; stream } code_batch =

{

task with

context = { ctx = task.context; stream };

})) )

let from_host (context : context) (tn : Tnode.t) =

{

context_lifetime = context;

description =

work;

});

true

{

context_lifetime = context;

description =

work;

});

true

Backend.to_buffer tn ~dst:merge_ptr ~src:src.ctx

in

let description =

in

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

in

let num_devices () = 1

let suggested_num_streams CPU = Domain.recommended_domain_count () - 1

let cleanup_stream stream =

assert (Domain.is_main_domain ());

await stream;

stream.state.keep_spinning <- false;

Stdlib.Condition.broadcast stream.state.dev_wait_for_work;

Domain.join stream.domain

let get_device ~ordinal =

CPU

let new_stream CPU =

Stdlib.Gc.finalise cleanup_stream stream;

stream

let get_stream_device _stream = CPU

let get_ctx_stream { stream; _ } = stream

let to_ordinal _ = 0

end

let will_wait_for _context () = ()

type stream = {

merge_buffer : (buffer_ptr * Tnode.t) option ref;

mutable allocated_buffer : (buffer_ptr * int) option;

}

Array.map routines

~f:(Option.map ~f:(fun task -> { task with context = { ctx = task.context; stream } })) )

let from_host (context : context) (tn : Tnode.t) =

Option.value ~default:false

let num_devices () = 1

let suggested_num_streams _device = !sync_suggested_num_streams

let get_device ~ordinal =

if ordinal <> 0 then invalid_arg "Sync_backend backends only have device number 0";

CPU

let new_stream CPU =

in

result

let get_stream_device _stream = CPU

let get_ctx_stream { stream; _ } = stream

let to_ordinal _ = 0

let lower_assignments ?name bindings asgns =

primary_context : Cu.Context.t;

mutable copy_merge_buffer : buffer_ptr;

mutable copy_merge_buffer_capacity : int;

released : Utils.atomic_bool;

cross_stream_candidates : ctx_array Hashtbl.M(Tn).t;

(** Freshly created arrays that might be shared across streams. The map can both grow and

(** The streams owning the given nodes. This map can only grow. *)

}

[@@deriving sexp_of]

and stream = {

device : device;

cu_stream : Cu.Stream.t;

mutable merge_buffer : (buffer_ptr * Tn.t) option;

}

(* We need this: there can be an arbitrary gap between the finalizer run and the deallocation. *)

if Atomic.get result.released then default () else result

(* Strange that we need ctx_set_current even with a single device! *)

set_ctx device.primary_context;

let cu_stream = Cu.Stream.create ~non_blocking:true () in

let cuda_properties =

let cache =

let get_ctx_stream { stream; _ } = stream

let get_stream_device { device; _ } = device

let to_ordinal { ordinal; _ } = ordinal

let await stream : unit =

set_ctx stream.device.primary_context;

~src:s_arr.ptr ~src_ctx:src.ctx dst.stream.cu_stream

in

if

then false

else

match Map.find src.ctx_arrays tn with

let data = Hashtbl.find_or_add device.cross_stream_candidates key ~default in

Map.add_exn ctx_arrays ~key ~data)

else if Tn.known_shared_cross_stream key then (

raise

@@ Utils.User_error

("Cuda_backend.alloc_if_needed: node " ^ Tn.debug_name key

^ " assumed to be cross-stream-shared but then written to on multiple devices")

let data = Hashtbl.find_exn device.cross_stream_candidates key in

Map.add_exn ctx_arrays ~key ~data)

else (

let get_ctx_stream Unimplemented_ctx = Unimplemented_stream

let get_name Unimplemented_stream : string = failwith "CUDA missing: install cudajit"

let to_ordinal _stream = 0

let name = "cuda"