arrayjit 0.4.0 · OCaml Package

Parameters

module Backend : Simple_backend

Signature

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 init : label:Base.string -> context

label is usually the backend name concatenated with the device number.

val finalize : context -> Base.unit

Finalizes (just) the context.

val alloc_buffer : 
  ?old_buffer:(buffer_ptr * Base.int) ->
  size_in_bytes:Base.int ->
  Base.unit ->
  buffer_ptr

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 link : 
  ?from_prior_context:Tnode.t Base.list ->
  merge_buffer:(buffer_ptr * Tnode.t) Base.option Base.ref ->
  context ->
  code ->
  routine

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

The from_prior_context nodes must not be added to the resulting context -- if needed in context, they must be part of the given context.

val link_batch : 
  ?from_prior_context:Tnode.t Base.list ->
  merge_buffer:(buffer_ptr * Tnode.t) Base.option Base.ref ->
  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 (in particular, the routines' contexts are not independent).

The from_prior_context nodes must not be added to the resulting context -- if needed in context, they must be part of the given context.

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

package arrayjit

Parameters

Signature