module Backend : No_device_backendinclude No_device_backendval sexp_of_code : code -> Sexplib0.Sexp.tval sexp_of_code_batch : code_batch -> Sexplib0.Sexp.tval sexp_of_buffer_ptr : buffer_ptr -> Sexplib0.Sexp.tval sexp_of_context : context -> Sexplib0.Sexp.ttype nonrec routine = context Backend_utils.Types.routineval sexp_of_routine : routine -> Sexplib0.Sexp.tval initialize : Backend_utils.Types.config -> Base.unitInitializes a backend before first use or (on some backends) after unsafe_cleanup. Does nothing if the backend is already initialized.
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 initialize before using the backend.
val finalize : context -> Base.unitFinalizes (just) the context.
val expected_merge_nodes : code_batch -> Tnode.t Base.option Base.arrayval compile :
?shared:Base.bool ->
?name:Base.string ->
Indexing.unit_bindings ->
Assignments.t ->
codeIf ~shared:true (default false), the backend should prefer to do more compile work in a device-agnostic way. If ~shared:false, the backend can opt to postpone compiling altogether until link is called, to benefit from more optimizations.
val compile_batch :
?shared:Base.bool ->
?names:Base.string Base.array ->
?occupancy:(name:Base.string -> src_n:Base.int -> Base.bool) ->
Indexing.unit_bindings ->
Assignments.t Base.array ->
code_batchUnlike the ~shared parameter, compile_batch vs. compile is mostly about improving the 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.
Cleans up all work on a backend, releases resources. All previously retrieved values (contexts, virtual and physical devices) become invalid. The backend needs to be initialized again to be used again.
val to_buffer : Tnode.t -> dst:buffer_ptr -> src:context -> Base.unitval host_to_buffer : Ndarray.t -> dst:buffer_ptr -> Base.unitval buffer_to_host : Ndarray.t -> src:buffer_ptr -> Base.unitval get_buffer : Tnode.t -> context -> buffer_ptr Base.optionReturns the routine for the code's procedure, in a new context derived from the given context.
val link_batch :
?from_prior_context:Tnode.t Base.list ->
context ->
code_batch ->
context * routine Base.option Base.arrayReturns the routines for the procedures included in the code batch. The returned context is downstream of all the returned routines.
If the array is both hosted and in-context, schedules a copy from host to context and returns true, otherwise returns false. NOTE: when run for a device, it's the caller's responsibility to synchronize the device before the host's data is overwritten.
If the array is both hosted and in-context, schedules a copy from context to host and returns true, otherwise returns false. NOTE: when run for a device, it's the caller's responsibility to synchronize the device before the host's data is read.
val device_to_device :
Tnode.t ->
into_merge_buffer:Backend_utils.Types.merge_buffer_use ->
dst:context ->
src:context ->
Base.boolIf the node is absent from the src context and either it is present in the dst context or ~into_merge_buffer is different from No: raises an error.
If ~into_merge_buffer:No: If the node is present in the dst context, schedules a copy of the tensor node from the device of src to the device of dst and returns true, otherwise returns false.
If ~into_merge_buffer is different from No: schedules the following task and returns true.
The merge-buffer task sets on dst the merge buffer source to the given node. If ~into_merge_buffer:Streaming, remembers the buffer pointer of the source node to use for streaming, without blocking. If ~into_merge_buffer:Copy, copies from src to the merge buffer of dst's device.
If the dst context resulted from a compilation with Streaming or Copy specific merge buffer code, the device_to_device call should fail immediately if there's a mismatch with ~into_merge_buffer.
NOTE: it's the caller's responsibility to synchronize the src device, if needed, before calling device_to_device, and if ~into_merge_buffer:Streaming, the dst device afterward, before any computations on the src device overwrite the node.
val alloc_buffer :
?old_buffer:(buffer_ptr * Base.int) ->
size_in_bytes:Base.int ->
device ->
buffer_ptrval await : device -> Base.unitBlocks till the device becomes idle, i.e. synchronizes the device.
val is_idle : device -> Base.boolWhether the device is currently waiting for work.
val sexp_of_device : device -> Base.Sexp.tval get_device : ordinal:Base.int -> physical_deviceval suggested_num_virtual_devices : physical_device -> Base.intThe optimal number of virtual devices for the given physical device to follow the Backend_types.config strategy passed to No_device_backend.initialize.
val new_virtual_device : physical_device -> deviceval get_physical_device : device -> physical_deviceval to_ordinal : physical_device -> Base.intval to_subordinal : device -> Base.intval get_name : device -> Base.string