github ↗
docs menu

Module Rig_edge

The contract between rig and the drivers that run its devices.

A driver is the library that runs one kind of hardware for rig. It allocates a device's memory, loads code on it and places work on its queues. A vendor library implements Driver and supplies an opener, which a program passes to Rig.open_. The driver, and its ABI library that declares the record its compiled code reads (capability), link this library and never rig. rig, which every program that opens a device links, defines the C function rig_code_sync that rig_edge.h declares. A library for a device whose memory the host reads and writes through calls, such as files, implements Io instead.

rig reads a device's facts once, when the device opens. It then calls the driver's OCaml functions to allocate and map memory, to load code and to grow its queues, and three C functions of the driver to run work (The C functions). room returns whether a submission's parts fit the device's queues, submit places them on the queues as the work of the next value of the device's timeline, and commit lets that work run.

A device reports progress in its timeline word, eight bytes of memory that hold a value of the device's timeline. Rig's preamble and Timelines section define values, reaching and committing. While the device is live, the word holds the last committed value v such that the work of every value up to v has completed. Only the driver writes the word, and the word never moves backwards. rig reads it to learn which values the device has reached, and other devices' queues may wait on it. After a loss the word may go on to show dropped values (Rig's Loss section).

The life of a device, as its driver sees it, goes like this:

  1. Rig.open_ calls the opener's make with the address of the device's word (opener) and reads the new device's facts.
  2. rig allocates, maps and frees memory, loads and unloads code, and hands work to the device through the C functions.
  3. rig loses the device for a Driver.Fault, a failed submit or commit, a hang, a close, a failure of the process, a fork, or the loss of a device whose value its queue waits for.
  4. rig calls Driver.stop once no call of the driver is in flight on the device (Driver). Once the word shows the device's last submitted value, the last value rig gave the driver's submit, rig frees the device's remaining memory and unloads its images, and the word may pass to a device that opens later. rig frees a word of the driver's own once no other device's work reads it (facts). rig keeps the device's handle while the program holds the device.

Each function's doc says what it does, and a driver's implementation must do that. Further obligations say "must".

While rig is 1.x, a minor release may change Driver, Io, the records they take and return, and rig_edge.h's structures. The types Rig re-exports, Rig.Buffer.kind and Rig.Buffer.dma, change only as Rig does. A driver outside rig must pin the exact rig version it builds against.

Facts

A device's facts are what rig reads of it once, when it opens (facts).

#type memory =
  1. #| Device
    (*

    The device's own memory.

    *)
  2. #| Pinned
    (*

    Page-locked host memory that the device's driver allocates for the device, which both the device's work and the host address. rig asks for none on a device that declares host_addresses.

    *)
  3. #| Mapped
    (*

    The device's own memory, which the host also addresses through a write-combined window.

    *)

The type for the kinds of a device's memory.

#type 'region mapping =
  1. #| Maps of 'region
    (*

    The region over the memory.

    *)
  2. #| Never
    (*

    The device does not map this memory, now or later, such as pages mapped read-only.

    *)
  3. #| No_room
    (*

    The device has no room for the mapping now, such as under a limit on locked memory or on mappings. rig may return memory and ask once more.

    *)

The type for a driver's answer to a request to map memory (Driver.map_host, Driver.map_peer).

#type kind =
  1. #| Words
    (*

    32-bit words placed on the queue.

    *)
  2. #| Fill
    (*

    A C function called with the queue's context.

    *)
  3. #| Copy
    (*

    A copy of rows of bytes between memory of the device, or a fill of its memory with a byte.

    *)
  4. #| Launch
    (*

    A function of an image run over a grid of groups of threads.

    *)

The type for the work of a part, as rig_edge.h tags it.

#type queue = {
  1. #name : string;
    (*

    The queue's name, such as "COMPUTE:0": its lane in a profile.

    *)
  2. #runs : kind list;
    (*

    The kinds of work the queue runs.

    *)
}

The type for a device's queues. rig places each part of a submission on one of them, and orders it after the parts of other queues it follows (The C functions):

  • A copy goes on the queue of the last earlier part of its submission that it follows, where that queue runs Copy, and otherwise on the device's copy engine, the queue its facts name in copy_engine (facts), where rig's own copies go.
  • Words go on the first queue whose runs lists both Launch and Words, and otherwise on the first that lists Words.
  • A launch or a fill goes on the first queue whose runs lists its kind.
#type completion =
  1. #| Store
    (*

    A store that the driver places on a queue writes the word.

    *)
  2. #| Host
    (*

    The host writes the word, from the driver's handler or before the driver's submit or commit returns (The C functions).

    *)

The type for how a device's word advances.

#type waits = {
  1. #stores : bool;
    (*

    Its queues wait for a Store producer.

    *)
  2. #hosts : bool;
    (*

    Its queues wait for a Host producer.

    *)
}

The type for the producers a device's queues wait for. A producer is a device whose value another device's work waits for. rig waits on the host for the producers a device's queues do not wait for.

#type capability =
  1. #| Capability : 'a Stdlib.Type.Id.t * 'a -> capability

The type for what compiled code needs from a device, a record that the driver's ABI library declares under its key.

#type 'region facts = {
  1. #arch : string;
    (*

    The device's architecture, such as "gfx1100".

    *)
  2. #budget : int;
    (*

    The bytes of own memory the device should hold at most.

    *)
  3. #queues : queue list;
    (*

    The device's queues.

    *)
  4. #copy_engine : int option;
    (*

    Some i if the device's copy engine is the queue at index i of queues, and None if it has none. A copy engine other than the first queue lets copies run beside the first queue's work.

    *)
  5. #completion : completion;
    (*

    How the word advances.

    *)
  6. #waits : waits;
    (*

    The producers its queues wait for.

    *)
  7. #may_block : bool;
    (*

    Whether room, submit and commit (The C functions) may block on the device's own earlier work, its own transfers or its library's back-pressure. Where they may, rig counts them (Driver), so Driver.stop never runs while one blocks.

    *)
  8. #hang_ms : int option;
    (*

    Some n if rig loses the device once its word stays below a committed value for n milliseconds, for a device whose work nothing else bounds, and None where the vendor's driver bounds it. n must be at least 1. The clock starts at a wait, so the loss comes n ms or more after the work stalled. Rig's Hangs section states when the clock starts again. rig passes the loss's reason, "no progress for n ms", to Driver.stop.

    *)
  9. #maps_host : bool;
    (*

    Whether the device maps host memory of its machine that starts on a page (Driver.map_host).

    *)
  10. #host_addresses : bool;
    (*

    Whether the host addresses every region the device allocates, with a host address in each one's location (Driver.locate). The host then reaches the device's memory and copies it, and rig asks for none of the device's Pinned memory.

    *)
  11. #alignment : int;
    (*

    The power of two that the address of every region Driver.alloc makes is a multiple of, for every memory kind: its address and its host address (Driver.locate).

    *)
  12. #cache_bytes : int option;
    (*

    Some n with n the bytes of the device's last-level cache (Rig.cache_bytes), and None where the driver does not know it. n must be at least 1.

    *)
  13. #counters : (string * int) list;
    (*

    The counters the device measures around each launch of a value that submit timed, each with how many values it has, at least 1, in the order measures writes them (The C functions). It is [] iff the device measures none, whose measures is NULL.

    *)
  14. #capability : capability;
    (*

    The ABI record, filled when the device opened.

    *)
  15. #word : 'region;
    (*

    The timeline word: the word make was given (opener), whose location's host address is that word's. Other devices' work may read it through rig's array. A device whose work cannot write that word, such as a device of another machine, has a word of the driver's own instead. rig reads a word with another host address, or none, through Driver.signaled, and raises the given word to each value it reads. A word whose host address is not the one rig gave counts as the driver's own, also one the opener's caller gave make. rig frees a word of the driver's own (Driver.free) only after Driver.stop, and never frees the word it gave.

    *)
  16. #bus : string option;
    (*

    The bus address of the device's PCI function on its machine, such as "0000:03:00.0", or None for a device that has none, such as a Metal device or a device of another machine.

    *)
  17. #edge : nativeint;
    (*

    The address of the device's C state, valid until its stop returns (Driver.stop). The state's first member must point to the driver's struct rig_driver of rig_edge.h (The C functions).

    *)
}

The type for the facts of a device. A Driver.Fault while rig reads them makes the open return Error. The open raises Invalid_argument if copy_engine names no queue of queues or one whose runs lacks Copy, hang_ms or cache_bytes is below 1, alignment is not a power of two, edge is 0n, or the state's first member is NULL. rig checks nothing else of edge. A region Driver.alloc makes that breaks host_addresses or alignment is a fault of the device: rig frees the region and loses the device, and the allocation raises Rig.Lost.

The C functions

A driver's struct rig_driver holds five C functions, which rig calls with the facts' edge. They return the constants of rig_edge.h.

  • room is given the parts and the number of waits the submit after it is given. It returns RIG_FITS if the parts and the waits fit the device's queues now, and RIG_LATER if they fit once one of its values is reached. It returns RIG_GROW for parts and waits that exceed the device's queues as they are but fit once they grow (Driver.grow), and RIG_NEVER for parts and waits that exceed the device's empty queues however they grow, for parts that name work it does not run or have kind 0, and for a launch whose grid or group is empty or exceeds the device's or the function's limits, or whose shared memory exceeds the function's.
  • submit encodes the parts on the device's queues as the work of the value it is given, after the waits it is given and the device's earlier work. Its calls are given the values b + 1, b + 2, … in order, where b is the value the word given to make holds when make returns. Each part names its queue, by its index in queues, and the parts of other queues it starts after (struct rig_part in rig_edge.h). When it returns, all of the value's work is on the queues, though the device may hold it until it is committed. It returns RIG_COMMITTED if every value up to its own is committed, RIG_OK if its value is only encoded, and RIG_FAILED if it failed, which loses the device. While a profile is taken, it is given a pair of times, where it must write when the value's work started and ended on the device, on the host's clock (rig_now_ns in rig_edge.h), or leave both 0 for work it cannot time. It writes them before the word shows the value if settle is NULL, and before settle of the value returns otherwise. rig gives the pair to no other value until then.
  • commit is given a value v no greater than the last value given to submit, and lets the work up to v run. Once that work has completed, the device writes v or a later value into the word. Committing a committed value must do nothing. It returns RIG_OK, or RIG_FAILED, which loses the device.
  • settle is given a value v that submit timed, once the word shows v, and must write v's times if it has not. rig calls it before it reads them, or before it gives their pair to another value. It may call it more than once for a value, one call at a time, and a later call finds the times written and leaves them. A value of a device lost first may get no call. A driver whose device's work writes the word, with no host code between the work's end and the word, has the work stamp its start and end in the driver's own memory, and converts the stamps to the host's clock in settle, or in a later submit once v's work has completed. settle never waits for the device's work. It may read a clock of the device, with one call of the operating system or a read of the device's registers. It returns RIG_OK, or RIG_FAILED if it could not write the times, such as for a stamp the work never wrote, which loses the device. settle may be NULL for a driver that writes the times itself before the word shows the value.

    The times are as exact as the driver's correlation of the device's clock with the host's, which its interface states. A driver converts in value order, and never puts a value's start before the end of the value converted before it, since a value's work starts once every earlier value's work has completed (device order).

  • measures is given a value v that submit timed and that has at least one launch, and is NULL for a driver whose facts list no counters. rig names no value without a launch, so a driver keeps nothing for one. Given an out of zeroes, once the word shows v, it writes the measures of v's launches, in the order of its parts, one row of 3 + n words each, with n the sum of counters' values: a status, 1 if measured and 0 if not, its start and its end on the host's clock, then each counter's values in the order of counters. It writes nothing for a launch it did not measure. Given NULL, it frees what it kept of v, at once or once the word shows v, and later calls do not name v. rig calls it with an out at each read of v's measures, and with NULL once, after the last. It returns RIG_OK, or RIG_FAILED, which loses the device.

rig calls the five one at a time for each device, with room before each submit. They must call no function of the OCaml runtime and read no OCaml value, since rig may call them with the runtime released. A submit, commit, settle or measures that returns RIG_FAILED sets *failure to the driver's message, which must live as long as the device.

Each submission holds RIG_SCRATCH bytes of scratch for each of its parts, which rig hands to room and submit with the parts. Either may write the scratch and read it back during the call, such as to keep which of the device's engines runs each part, so that neither allocates. No other call uses the scratch meanwhile, and a call finds what an earlier call left there. A copy part's fields may differ from one call to the next, since a copy's ends may be a run's buffers (Rig.Submission.side). A driver must decide again on each call what it derives from them.

A driver whose submit always returns RIG_COMMITTED is never asked to commit. Every other driver must commit on its own at least once every k values it receives, so that the word shows each value at most k values late while work is submitted. k is the driver's own constant, which rig does not read, and the driver's interface states it.

A driver must keep three properties:

  • Prefix order: the work handed to each queue runs in the order it was handed.
  • Device order: a value's work starts after the previous value's work has completed, on every queue. The driver's submit meets it by making its work wait for the queue that completed the previous value.
  • Observable completion: the host learns, through the word, that a prefix of the values has completed, or learns that the device faulted.

Memory

#type location = {
  1. #address : int option;
    (*

    The region's address as the device's work addresses it, or None for memory named by its handle only.

    *)
  2. #host : int option;
    (*

    The host address of the region's first byte, if the host addresses it.

    *)
  3. #handle : nativeint;
    (*

    The driver's object for the region, such as a buffer object, a CUdeviceptr or an MTLBuffer.

    *)
}

The type for where a region lies.

#type dma =
  1. #| Bus of (int * int) list
    (*

    The memory's bus addresses on the device's machine, as (address, bytes) pairs in order, each a multiple of 4096. They hold until the process exits.

    *)
  2. #| Dmabuf of {
    1. #fd : int;
    2. #offset : int;
    }
    (*

    A new dma-buf file descriptor, close-on-exec, which the caller owns and closes, and the offset of the region's first byte in it. The dma-buf keeps the memory until its last reference goes, such as a NIC's registration of it.

    *)

The type for how another PCI function of a device's machine, such as a NIC, reaches a region's memory by DMA.

Code

#type ('region, 'image) code =
  1. #| Loaded of 'image
    (*

    The driver's library placed the code itself.

    *)
  2. #| Place of int * 'region -> 'image * string
    (*

    The code goes into n bytes of the device's Device memory. rig allocates a region of at least n bytes and calls the function with it, and the function returns the image and the bytes, at most n, that rig places at the region's start. The driver must make nothing on the device before the function is called, so a Place whose function is never called leaves nothing to release. The function must call nothing that may block or fault. rig calls it as a counted call (Drivers).

    *)

The type for a binary a driver took.

#type entry = {
  1. #code : int;
    (*

    What compiled code names the function by, such as a kernel descriptor's address, a code address, a CUfunction or a pipeline.

    *)
  2. #launch : nativeint;
    (*

    What the driver's submit reads to launch the function (launch.launch in rig_edge.h), such as its dispatch template and limits, valid while the image is loaded; 0n for a driver whose queues run no Launch.

    *)
  3. #params : int option;
    (*

    The number of parameter bytes, from byte 0, that the function reads, or None if the driver does not know it. rig refuses a launch that gives fewer, or that places a ref past them.

    *)
  4. #threads : int option;
    (*

    The most threads a group of the function may have, or None if the driver does not know it. rig refuses a run that sets more.

    *)
}

The type for a function of an image, as a driver names it.

#type entry_answer =
  1. #| Entry of entry
    (*

    The function.

    *)
  2. #| Absent
    (*

    The image has no function of that name.

    *)
  3. #| No_memory of int
    (*

    The device has no memory for the bytes the function's launches need, such as its scratch. rig reclaims memory and asks again (Driver.entry).

    *)

The type for a driver's answer when asked for a function of an image.

#type grown =
  1. #| Grown
    (*

    The growth is made; the next room puts it into the queues.

    *)
  2. #| Short of int
    (*

    The device has no memory for the bytes the growth takes. rig reclaims memory and asks again.

    *)

The type for a driver's answer when asked to grow a device's queues (Driver.grow).

Drivers

#module type Driver = sig ... end

Drivers.

#module type Io = sig ... end

Devices of memory that the host reads and writes through calls.

Opening

#type 'a opener = {
  1. #name : string;
    (*

    The device's name, without a machine suffix.

    *)
  2. #driver : (module Driver with type t = 'a);
    (*

    The device's driver.

    *)
  3. #make : word:nativeint -> ('a, string) Stdlib.result;
    (*

    Opens a new device handle, whose timeline word lies at the host address word.

    *)
}

The type for openers of named devices. Making an opener does no I/O and opens nothing. make ~word opens the device, and returns Error when the device cannot be opened, such as when the hardware is missing. Rig.open_ calls make only when no device of name is open on its machine. It gives it the address of the device's word in rig's array of timeline words (rig_words in rig_edge.h), and the device's values start after the value the word holds. make may raise that value, never lower it, and an Error leaves no work that writes the word, which rig then gives to a later device. A success hands the handle to rig, which stops it (Driver.stop) at the device's close, or at an open of its name once it is lost. A caller that calls make itself owns the handle and must stop it, and gives it 8 bytes of host memory, aligned to 8, that live as long as the handle. Each driver's opener states whether make opens another handle while an earlier one is not stopped. A make that may take a lock must raise Invalid_argument, having opened nothing, while its thread holds one (Drivers).

#type 'a io_opener = {
  1. #name : string;
    (*

    The device's name, without a machine suffix.

    *)
  2. #io : (module Io with type t = 'a);
    (*

    The device's io library.

    *)
  3. #make : unit -> ('a, string) Stdlib.result;
    (*

    Opens a new io device handle.

    *)
}

The type for openers of named io devices. Rig.open_io calls make and owns its handles as Rig.open_ does for an opener. A caller that calls make itself owns the handle and must call its stop (Io).