Module `Tolk.Device`

Device runtime abstraction.

A device bundles the pieces needed to run compiled kernels on a specific backend: an Allocator.packed for buffer management, a Renderer_set.t for renderer/compiler selection, a Queue.t for kernel dispatch, and a preparation hook for device-specific program setup.

Buffer.t values are existentially packed so that the concrete backend buffer type does not leak into consumer code. Program.t values carry compiled binaries together with their runtime metadata. Compiled programs are cached per device and compiler context.

Types

type t

The type for compiled device runtimes.

type device = t

Alias for t, used in signatures where device reads better than Device.t.

Buffer specification

module Buffer_spec : sig ... end

Buffer allocation options.

Allocator

module Allocator : sig ... end

Backend allocator interface.

LRU allocator

module Lru_allocator : sig ... end

LRU buffer reuse layer.

Buffers

module Buffer : sig ... end

Existentially-packed device buffers.

Runtime program handle

type prog = {

call : nativeint array -> global:int array -> local:int array option -> vals:int64 array -> wait:bool -> timeout:int option -> float option;
free : unit -> unit;
handle : nativeint;
(*
Backend kernel handle used to build Graph nodes. 0n when the backend has no addressable kernel object.
*)

}

A device-specific dispatch handle.

type runtime = string -> bytes -> runtimevars:(string * int) list -> prog

runtime name lib ~runtimevars creates a dispatch handle for lib with entry point name. runtimevars maps variable names (e.g. "core_id") to their index in the vals array.

Batched dispatch graphs

module Graph : sig ... end

Backend interface for batched replay of a fixed call sequence.

Renderer selection

module Renderer_set : sig ... end

Available renderers for a device.

Device operations

val make : 
  name:string ->
  allocator:Allocator.packed ->
  renderer_set:Renderer_set.t ->
  runtime:runtime ->
  synchronize:(unit -> unit) ->
  ?invalidate_caches:(unit -> unit) ->
  ?graph:Graph.t ->
  unit ->
  t

make ~name ~allocator ~renderer_set ~runtime ~synchronize ?invalidate_caches ?graph () is a device runtime.

runtime name lib loads a compiled binary and returns a dispatch handle.

synchronize () blocks until all pending work on the device completes.

graph is the batched-dispatch capability, or absent when the backend cannot replay call sequences as a single dispatch.

val name : t -> string

name d is d's device name.

val renderer : t -> Renderer.t

renderer d is the active renderer.

val runtime : t -> runtime

runtime d is d's runtime factory.

val synchronize : t -> unit

synchronize d blocks until all pending work on d completes.

val graph : t -> Graph.t option

graph d is d's batched-dispatch capability, if any.

val compile_program : 
  t ->
  ?name:string ->
  ?applied_opts:Tolk_uop.Uop.Opt.t list ->
  ?estimates:Program_spec.Estimates.t ->
  Program_spec.program ->
  Program_spec.t

compile_program d ?name ?estimates program renders and compiles program for d, returning a prepared Program.t.

Results are cached by device name, compiler name, kernel content digest, renderer context, entry name, and estimates. Cached programs are cloned (entry address and cleanup cleared) before being returned.

name defaults to "kern". estimates defaults to Program_spec.Estimates.zero.

val create_buffer : 
  size:int ->
  dtype:Tolk_uop.Dtype.t ->
  ?spec:Buffer_spec.t ->
  t ->
  Buffer.t

create_buffer ~size ~dtype ?spec d is an unallocated buffer for size elements of dtype on d.

spec defaults to Buffer_spec.default.

val invalidate_caches : t -> unit

invalidate_caches d flushes device caches (e.g., L2) if the device supports it. No-op if ~invalidate_caches was not provided to make. Called by beam search between timing runs for consistent measurements.

Device registry

The registry maps canonical device names to opened device runtimes. Backends register an opener per name prefix (e.g. "CPU"); get opens a device on first lookup and caches it, so every consumer of a device name shares one runtime instance per canonical name.

val canonicalize : string -> string

canonicalize name is name with its backend prefix uppercased and a trailing ":0" instance suffix removed, so "cpu:0", "CPU:0" and "CPU" all name the same device.

val register : string -> (string -> t) -> unit

register prefix opener installs opener for device names whose backend prefix is prefix (case-insensitive). opener name must return the device runtime for the canonical name.

val get : string -> t

get name is the device runtime for the canonicalized name, opened via its registered opener on first lookup and cached afterwards.

Raises Failure if no opener is registered for name's prefix or the opener fails.

Multi-device buffers

module Multi_buffer : sig ... end

Buffers spanning multiple devices.