Module Context.Backends_deprecated

A collection of the execution backends

module Schedulers : sig ... end
val plan_pool_segments : cap:Base.int -> what:Base.string -> debug_name:(Base.int -> Base.string) -> (Base.int * Base.int) Base.list -> (Base.int * Base.int) Base.list * Base.int Base.list

gh-ocannl-344 pool-allocator planner. Lays out (size, alignment) allocations (in order) into pools so no pool's bumped extent exceeds cap bytes (the uint32 4 GB per-pool ceiling when large_models = false). Returns each item's (segment_index, byte_offset) and the byte size of each segment. Raises Ir.Utils.User_error (naming what and debug_name i) when a single item exceeds cap. Exposed for unit testing the segmenting/cap behavior with synthetic sizes.

val finalize : 'dev 'runner 'event. (module Ir.Backend_intf.Backend with type dev = 'dev and type event = 'event and type runner = 'runner) -> ('dev, 'runner, 'event) Ir.Backend_intf.context -> Base.unit

Frees the pools that are specific to the context -- not contained in the parent context. Note: use finalize to optimize memory, it is not obligatory because all pools are freed when their backend buffers are garbage-collected.

The implemented backends

Each backend is instantiated once per process, so its context type is nameable and two independently-created contexts on the same backend unify -- the precondition for Context.copy dispatching to the backend's device_to_device via wrapped_context. Instantiation touches no driver or hardware: device discovery stays lazy inside the backends (forced at first get_device, where Context.auto's fallback can catch an unusable driver/device per call), and on platforms without the corresponding library the dune-selected missing stub is what gets instantiated -- harmless at init, raising on use. Backend caches consequently persist across Tensor.unsafe_reinitialize; that is safe because tnode identity (Tnode.uid) is never reused.

type backend =
  1. | Cc
  2. | Multidev_cc
  3. | Cuda
  4. | Hip
  5. | Metal

The implemented backends. Constructors statically imply the corresponding singleton module (Cc -> Cc_b, ...).

val sexp_of_backend : backend -> Sexplib0.Sexp.t
val backend_of_sexp : Sexplib0.Sexp.t -> backend
val equal_backend : backend -> backend -> Base.bool
val get_backend : ?backend_name:Base.string -> Base.unit -> backend

The backend corresponding to backend_name, or if omitted, selected via the global backend setting. Non-generative: the same backend value (hence the same singleton state) each call.

val backend_name : backend -> Base.string

Inverse of get_backend's name parsing.

val backend_module : backend -> (module Ir.Backend_intf.Backend)

The singleton module as an existentially-packed first-class module, for generic consumers that thread a single backend through (the raw-API tests, Parallel). Code that must re-correlate two contexts later (e.g. Context.copy) should use wrapped_context instead: this projection erases the type components.

Contexts wrapped with their backend

A closed disjunction over the implemented backends' context types: matching two values on the same constructor recovers type equality directly, which is what lets Context.copy fall onto the backend-specific device_to_device when both contexts come from the same backend.

type wrapped_context =
  1. | Cc_ctx of Cc_b.context
  2. | Multidev_cc_ctx of Multidev_cc_b.context
  3. | Cuda_ctx of Cuda_b.context
  4. | Hip_ctx of Hip_b.context
  5. | Metal_ctx of Metal_b.context
val wrapped_backend : wrapped_context -> backend
val make_context : ?device_id:Base.int -> backend -> wrapped_context

A fresh root context (empty optimize_ctx) on the backend's device device_id (default 0). Raises when the backend's hardware or library is unavailable.

type 'a ctx_op = {
  1. f : 'dev 'runner 'event. (module Ir.Backend_intf.Backend with type dev = 'dev and type event = 'event and type runner = 'runner) -> ('dev, 'runner, 'event) Ir.Backend_intf.context -> ('dev, 'runner, 'event) Ir.Backend_intf.context * 'a;
}

A context-transforming backend operation, polymorphic over the backend's type components so with_backend can rebuild the same wrapped_context constructor around the result.

val with_backend : wrapped_context -> 'a ctx_op -> wrapped_context * 'a
type 'a ctx_query = {
  1. q : 'dev 'runner 'event. (module Ir.Backend_intf.Backend with type dev = 'dev and type event = 'event and type runner = 'runner) -> ('dev, 'runner, 'event) Ir.Backend_intf.context -> 'a;
}

A read-only backend operation; like ctx_op but leaves the context untouched.

val query : wrapped_context -> 'a ctx_query -> 'a