Module Schedulers.Multidev

Parameters

Signature

include Ir.Backend_intf.Backend_device_common
include Ir.Backend_intf.Device
include Ir.Backend_intf.Device_types
include Ir.Backend_intf.Device_config_common
type dev

Interface to a device driver.

val sexp_of_dev : dev -> Sexplib0.Sexp.t
type runner

Interface to a stream driver.

val sexp_of_runner : runner -> Sexplib0.Sexp.t
type event

An event tracks if a device's runner finished computing past a particular point in its schedule. These values are used internally for scheduling across devices/queues of the backend, and can be used for explicit scheduling.

val sexp_of_event : event -> Sexplib0.Sexp.t
val name : Base.string
type nonrec device = (dev, runner, event) Ir.Backend_intf.device
val sexp_of_device : device -> Sexplib0.Sexp.t
type nonrec context = (dev, runner, event) Ir.Backend_intf.context
val sexp_of_context : context -> Sexplib0.Sexp.t
include Ir.Backend_intf.Slab_alloc with type device := device
val alloc_pool : ?mode:Ir.Tnode.memory_mode -> device -> pool_id:Base.int -> size_in_bytes:Base.int -> alignment:Base.int -> Base.unit

Allocates the slab for pool_id on device. The optional ?mode carries the tnode's memory mode so backends can pick a storage mode (Metal private vs. shared); backends that do not care ignore it.

val free_pool : (device -> pool_id:Base.int -> Base.unit) Base.option

Frees the slab for pool_id and drops its table entry. None for backends that rely on GC.

val memset_zero : device -> pool_id:Base.int -> offset:Base.int -> size_in_bytes:Base.int -> Base.unit

Zero-initializes size_in_bytes at base_of pool_id + offset.

val make_device : dev -> runner -> ordinal:Base.int -> device
val make_context : ?ctx_buffers:Ir.Backend_intf.ctx_buffers -> ?optimize_ctx:Ir.Low_level.optimize_ctx -> device -> context

Returns a context without a parent.

val make_child : ?ctx_buffers:Ir.Backend_intf.ctx_buffers -> ?optimize_ctx:Ir.Low_level.optimize_ctx -> ?merge_buffer_node:Ir.Tnode.t Base.option -> context -> context

Returns a context with the same Backend_intf.context.device, Backend_intf.context.ctx_buffers, Backend_intf.context.optimize_ctx, Backend_intf.context.merge_buffer_node if omitted, as the given context's, which is also the Backend_intf.context.parent.

val get_name : device -> Base.string
val sync : event -> Base.unit

Blocks till the event completes, if it's not done already.

It is rarely needed to call sync explicitly, because it should always be called internally when necessary, in particular before extracting values from host.

val is_done : event -> Base.bool

Whether the event completed.

val will_wait_for : context -> event -> Base.unit

Schedules waiting for the given event on the context's device.

NOTE: it should rarely be needed to call will_wait_for explicitly, because it should always be called internally when necessary.

val static_properties : Base.unit -> Base.Sexp.t

Returns a sexp description of the properties of all devices. A function so that computing it (device enumeration) does not run at backend-module initialization: singleton backends instantiate eagerly at program startup, where touching a driver could fail runs that never use the backend.

val hardware_limits : Base.unit -> Ir.Backend_intf.hardware_limits

Conservative per-workgroup device limits: on multi-device backends the minimum across the devices, so code compiled once (compilation is not per-device) is valid wherever it links. All-None for backends that do not bind hardware axes. A function for the same reason as static_properties: computing it (device enumeration) must not run at backend-module initialization.

val get_used_memory : device -> Base.int

Returns (an upper bound of) the memory used for arrays, in bytes.

val get_global_debug_info : Base.unit -> Base.Sexp.t

Global debug information; backend-specific and might evolve independently on the backends.

val get_debug_info : device -> Base.Sexp.t

Per-device debug information; backend-specific and might evolve independently on the backends

val await : device -> Base.unit

Blocks till the device becomes idle, i.e. synchronizes the device's runner.

val all_work : device -> event

Returns the event indicating if any currently running or scheduled computations on the device have completed.

val is_idle : device -> Base.bool

Whether the device's runner is currently waiting for work.

val get_device : ordinal:Base.int -> device
val num_devices : Base.unit -> Base.int
include Ir.Backend_impl.Buffer with type buffer_ptr = Backend.buffer_ptr
type buffer_ptr = Backend.buffer_ptr
val sexp_of_buffer_ptr : buffer_ptr -> Sexplib0.Sexp.t

Backend-private buffer_loc -> base resolution, used by the backend's own transfer/link implementations. Not part of Backend_intf.

val schedule_task : device -> Ir.Task.t -> Base.unit