Module `Device.Buffer`

Existentially-packed device buffers.

A buffer is either a base buffer (directly allocated) or a view into a base buffer at a byte offset. Views share the base buffer's backing storage.

Buffers start unallocated. Call allocate or ensure_allocated to materialise backing storage. Each buffer has a globally unique id assigned at creation. A GC finaliser calls deallocate when the buffer becomes unreachable.

Reference counting (uop_refcount, add_ref) is managed externally by the compiler runtime and is not used for deallocation.

Types

type t

The type for existentially-packed device buffers.

Constructors

val create : 
  device:string ->
  size:int ->
  dtype:Tolk_uop.Dtype.t ->
  ?spec:Buffer_spec.t ->
  Allocator.packed ->
  t

create ~device ~size ~dtype ?spec allocator is an unallocated base buffer for size elements of dtype on device.

spec defaults to Buffer_spec.default.

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

view b ~size ~dtype ~offset is a view into b starting at byte offset and spanning size elements of dtype. The view shares the base buffer's allocator and spec.

Raises Invalid_argument if offset is negative, >= nbytes b, or if the resulting view extends past the root base buffer.

Identity and metadata

val id : t -> int

id b is b's globally unique identifier.

val base_id : t -> int

base_id b is the unique identifier of b's root base buffer. Equal to id b when b is itself a base buffer.

val device : t -> string

device b is the device name b is bound to.

val size : t -> int

size b is the element count.

val dtype : t -> Tolk_uop.Dtype.t

dtype b is the element dtype.

val spec : t -> Buffer_spec.t

spec b is the buffer specification.

val nbytes : t -> int

nbytes b is the size in bytes (size b * Dtype.itemsize (dtype b)).

val base : t -> t

base b is the root base buffer. If b is already a base buffer, base b is b itself.

val offset : t -> int

offset b is the byte offset into the base buffer. 0 for base buffers.

Allocation

val allocate : t -> unit

allocate b materialises backing storage for b. For views, ensures the base buffer is allocated first, then creates the offset view via the allocator.

Raises Invalid_argument if b is already allocated, or if b is a view and the allocator does not support Allocator.t.offset.

val ensure_allocated : t -> unit

ensure_allocated b calls allocate if b is not yet initialised. No-op otherwise.

val is_allocated : t -> bool

is_allocated b is true iff the base buffer's backing storage exists.

val is_initialized : t -> bool

is_initialized b is true iff this specific buffer or view has its own storage pointer set. A view can be uninitialised even when the base buffer is allocated.

val deallocate : t -> unit

deallocate b releases backing storage if allocated. For base buffers, frees via the allocator. For views, detaches from the base buffer. No-op if already deallocated.

Raises Invalid_argument if b is a base buffer that still has allocated views.

val supports_offset : t -> bool

supports_offset b is true iff b's allocator provides offset views.

val supports_transfer : t -> t -> bool

supports_transfer dst src is true iff dst's allocator provides native transfer and dst and src are on the same backend prefix (for example "METAL" for "METAL:0" and "METAL:1").

val allocator : t -> Allocator.packed

allocator b is the allocator of b's base buffer.

Reference counting

val uop_refcount : t -> int

uop_refcount b is the base buffer's UOp reference count.

val add_ref : t -> int -> t

add_ref b cnt increments the base buffer's UOp reference count by cnt and returns b.

Data transfer

val copyin : t -> bytes -> unit

copyin b src copies src into b.

Raises Invalid_argument if Bytes.length src <> nbytes b or if b is not allocated.

val copyout : t -> bytes -> unit

copyout b dst copies the contents of b into dst.

Raises Invalid_argument if Bytes.length dst <> nbytes b or if b is not allocated.

val as_bytes : t -> bytes

as_bytes b is a fresh bytes value containing the contents of b. Equivalent to allocating Bytes.create (nbytes b) and calling copyout.

val transfer : dst:t -> src:t -> bool

transfer ~dst ~src copies src into dst via dst's allocator transfer hook when supports_transfer is true. Returns true when native transfer was used and false when no transfer hook is available. Both buffers are allocated if transfer is used.

Raises Invalid_argument if dst and src differ in size or dtype.

val copy_between : dst:t -> src:t -> unit

copy_between ~dst ~src copies the contents of src into dst via a host-memory bounce buffer. Both buffers are allocated if needed.

Raises Invalid_argument if size dst <> size src or dtype dst <> dtype src.

val addr : t -> nativeint

addr b is the device address of b. Allocates b if needed.