package arrayjit

  1. Overview
  2. Docs

Module type Backends.BackendSource

include No_device_backend
type code
val sexp_of_code : code -> Sexplib0.Sexp.t
type code_batch
val sexp_of_code_batch : code_batch -> Sexplib0.Sexp.t
type buffer_ptr
val sexp_of_buffer_ptr : buffer_ptr -> Sexplib0.Sexp.t
type context
val sexp_of_context : context -> Sexplib0.Sexp.t
type nonrec routine = context Backend_utils.Types.routine
val sexp_of_routine : routine -> Sexplib0.Sexp.t
val name : Base.string

Initializes a backend before first use or (on some backends) after unsafe_cleanup. Does nothing if the backend is already initialized.

val is_initialized : Base.unit -> Base.bool

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

Finalizes (just) the context.

val expected_merge_node : code -> Tnode.t Base.option
val expected_merge_nodes : code_batch -> Tnode.t Base.option Base.array
val compile : ?shared:Base.bool -> ?name:Base.string -> Indexing.unit_bindings -> Assignments.t -> code

If ~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_batch

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

val unsafe_cleanup : Base.unit -> Base.unit

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.unit
val host_to_buffer : Ndarray.t -> dst:buffer_ptr -> Base.unit
val buffer_to_host : Ndarray.t -> src:buffer_ptr -> Base.unit
val get_buffer : Tnode.t -> context -> buffer_ptr Base.option

Returns the routine for the code's procedure, in a new context derived from the given context.

Returns the routines for the procedures included in the code batch. The returned context is downstream of all the returned routines.

val from_host : context -> Tnode.t -> Base.bool

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.

val to_host : context -> Tnode.t -> Base.bool

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

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

type physical_device
type device
val init : device -> context
val alloc_buffer : ?old_buffer:(buffer_ptr * Base.int) -> size_in_bytes:Base.int -> device -> buffer_ptr
val await : device -> Base.unit

Blocks till the device becomes idle, i.e. synchronizes the device.

val is_idle : device -> Base.bool

Whether the device is currently waiting for work.

val sexp_of_device : device -> Base.Sexp.t
val get_device : ordinal:Base.int -> physical_device
val num_physical_devices : Base.unit -> Base.int
val suggested_num_virtual_devices : physical_device -> Base.int

The 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 -> device
val get_ctx_device : context -> device
val get_physical_device : device -> physical_device
val to_ordinal : physical_device -> Base.int
val to_subordinal : device -> Base.int
val get_name : device -> Base.string