Module `Device.Multi_buffer`

Buffers spanning multiple devices.

A multi-device buffer holds one Buffer.t per device, all sharing the same size and dtype. Operations apply element-wise across the per-device buffers.

Types

type t

The type for multi-device buffers.

Constructors

val create : 
  devices:string list ->
  size:int ->
  dtype:Tolk_uop.Dtype.t ->
  ?spec:Buffer_spec.t ->
  unit ->
  t

create ~devices ~size ~dtype ?spec () is a multi-device buffer with one underlying buffer per device name in devices, each resolved through the device registry (get).

spec defaults to Buffer_spec.default. The trailing unit argument is needed because spec is optional.

Raises Invalid_argument if devices is empty and Failure if a device name cannot be opened.

val of_bufs : Buffer.t list -> t

of_bufs bufs is a multi-device buffer stacking the per-device buffers bufs.

Raises Invalid_argument if bufs is empty or the buffers disagree in size or dtype.

val view : t -> size:int -> dtype:Tolk_uop.Dtype.t -> offset:int -> t

view t ~size ~dtype ~offset is a multi-device buffer viewing each underlying buffer at byte offset for size elements of dtype. See Buffer.view.

Accessors

val bufs : t -> Buffer.t list

bufs t is the underlying per-device buffers, one per device in the order given to create.

val size : t -> int

size t is the element count (same across all buffers).

val dtype : t -> Tolk_uop.Dtype.t

dtype t is the element dtype (same across all buffers).

val is_allocated : t -> bool

is_allocated t is true iff all underlying buffers are allocated.

Operations

val add_ref : t -> int -> t

add_ref t cnt increments the UOp reference count on all underlying buffers by cnt and returns t.