arrayjit 0.4.0 · OCaml Package

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

val initialize : Backend_utils.Types.config -> Base.unit

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

val link : ?from_prior_context:Tnode.t Base.list -> context -> code -> routine

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

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

package arrayjit