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:
Rig.open_calls the opener'smakewith the address of the device's word (opener) and reads the new device's facts.- rig allocates, maps and frees memory, loads and unloads code, and hands work to the device through the C functions.
- rig loses the device for a
Driver.Fault, a failedsubmitorcommit, a hang, a close, a failure of the process, a fork, or the loss of a device whose value its queue waits for. - rig calls
Driver.stoponce 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'ssubmit, 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 = - #
| Device(*The device's own memory.
*) - #
| 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. - #
| 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 = The type for a driver's answer to a request to map memory (Driver.map_host, Driver.map_peer).
type kind = The type for the work of a part, as rig_edge.h tags it.
type queue = {- #
name : string;(*The queue's name, such as
*)"COMPUTE:0": its lane in a profile. - #
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 incopy_engine(facts), where rig's own copies go. - Words go on the first queue whose
runslists bothLaunchandWords, and otherwise on the first that listsWords. - A launch or a fill goes on the first queue whose
runslists its kind.
type completion = - #
| Store(*A store that the driver places on a queue writes the word.
*) - #
| Host(*The host writes the word, from the driver's handler or before the driver's
*)submitorcommitreturns (The C functions).
The type for how a device's word advances.
type waits = {- #
stores : bool;(*Its queues wait for a
*)Storeproducer. - #
hosts : bool;(*Its queues wait for a
*)Hostproducer.
}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.
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 = {- #
arch : string;(*The device's architecture, such as
*)"gfx1100". - #
budget : int;(*The bytes of own memory the device should hold at most.
*) - #
queues : queue list;(*The device's queues.
*) - #
copy_engine : int option;(*
*)Some iif the device's copy engine is the queue at indexiofqueues, andNoneif it has none. A copy engine other than the first queue lets copies run beside the first queue's work. - #
completion : completion;(*How the word advances.
*) - #
waits : waits;(*The producers its queues wait for.
*) - #
may_block : bool;(*Whether
*)room,submitandcommit(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), soDriver.stopnever runs while one blocks. - #
hang_ms : int option;(*
*)Some nif rig loses the device once its word stays below a committed value fornmilliseconds, for a device whose work nothing else bounds, andNonewhere the vendor's driver bounds it.nmust be at least1. The clock starts at a wait, so the loss comesnms 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", toDriver.stop. - #
maps_host : bool; - #
host_addresses : bool;(*Whether the host addresses every region the device allocates, with a
*)hostaddress 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'sPinnedmemory. - #
alignment : int;(*The power of two that the address of every region
*)Driver.allocmakes is a multiple of, for every memory kind: itsaddressand itshostaddress (Driver.locate). - #
cache_bytes : int option;(*
*)Some nwithnthe bytes of the device's last-level cache (Rig.cache_bytes), andNonewhere the driver does not know it.nmust be at least1. - #
counters : (string * int) list;(*The counters the device measures around each launch of a value that
*)submittimed, each with how many values it has, at least1, in the ordermeasureswrites them (The C functions). It is[]iff the device measures none, whosemeasuresisNULL. - #
capability : capability;(*The ABI record, filled when the device opened.
*) - #
word : 'region;(*The timeline word: the word
*)makewas 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, throughDriver.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 gavemake. rig frees a word of the driver's own (Driver.free) only afterDriver.stop, and never frees the word it gave. - #
bus : string option;(*The bus address of the device's PCI function on its machine, such as
*)"0000:03:00.0", orNonefor a device that has none, such as a Metal device or a device of another machine. - #
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'sstruct rig_driverofrig_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.
roomis given the parts and the number of waits thesubmitafter it is given. It returnsRIG_FITSif the parts and the waits fit the device's queues now, andRIG_LATERif they fit once one of its values is reached. It returnsRIG_GROWfor parts and waits that exceed the device's queues as they are but fit once they grow (Driver.grow), andRIG_NEVERfor parts and waits that exceed the device's empty queues however they grow, for parts that name work it does not run or have kind0, 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.submitencodes 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 valuesb + 1,b + 2, … in order, wherebis the value the word given tomakeholds whenmakereturns. Each part names its queue, by its index inqueues, and the parts of other queues it starts after (struct rig_partinrig_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 returnsRIG_COMMITTEDif every value up to its own is committed,RIG_OKif its value is only encoded, andRIG_FAILEDif 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_nsinrig_edge.h), or leave both0for work it cannot time. It writes them before the word shows the value ifsettleisNULL, and beforesettleof the value returns otherwise. rig gives the pair to no other value until then.commitis given a valuevno greater than the last value given tosubmit, and lets the work up tovrun. Once that work has completed, the device writesvor a later value into the word. Committing a committed value must do nothing. It returnsRIG_OK, orRIG_FAILED, which loses the device.settleis given a valuevthatsubmittimed, once the word showsv, and must writev'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 insettle, or in a latersubmitoncev's work has completed.settlenever 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 returnsRIG_OK, orRIG_FAILEDif it could not write the times, such as for a stamp the work never wrote, which loses the device.settlemay beNULLfor 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).
measuresis given a valuevthatsubmittimed and that has at least one launch, and isNULLfor a driver whose facts list nocounters. rig names no value without a launch, so a driver keeps nothing for one. Given anoutof zeroes, once the word showsv, it writes the measures ofv's launches, in the order of its parts, one row of3 + nwords each, withnthe sum ofcounters' values: a status,1if measured and0if not, its start and its end on the host's clock, then each counter's values in the order ofcounters. It writes nothing for a launch it did not measure. GivenNULL, it frees what it kept ofv, at once or once the word showsv, and later calls do not namev. rig calls it with anoutat each read ofv's measures, and withNULLonce, after the last. It returnsRIG_OK, orRIG_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
submitmeets 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 = {- #
address : int option;(*The region's address as the device's work addresses it, or
*)Nonefor memory named by its handle only. - #
host : int option;(*The host address of the region's first byte, if the host addresses it.
*) - #
handle : nativeint;(*The driver's object for the region, such as a buffer object, a
*)CUdeviceptror anMTLBuffer.
}The type for where a region lies.
type dma = - #
| 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.
*) - #
| Dmabuf of {}(*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 = - #
| Loaded of 'image(*The driver's library placed the code itself.
*) - #
| Place of int * 'region -> 'image * string(*The code goes into
*)nbytes of the device'sDevicememory. rig allocates a region of at leastnbytes and calls the function with it, and the function returns the image and the bytes, at mostn, that rig places at the region's start. The driver must make nothing on the device before the function is called, so aPlacewhose 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 = {- #
code : int;(*What compiled code names the function by, such as a kernel descriptor's address, a code address, a
*)CUfunctionor a pipeline. - #
launch : nativeint;(*What the driver's
*)submitreads to launch the function (launch.launchinrig_edge.h), such as its dispatch template and limits, valid while the image is loaded;0nfor a driver whose queues run noLaunch. - #
params : int option;(*The number of parameter bytes, from byte
*)0, that the function reads, orNoneif the driver does not know it. rig refuses a launch that gives fewer, or that places a ref past them. - #
threads : int option;(*The most threads a group of the function may have, or
*)Noneif 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 = - #
| Entry of entry(*The function.
*) - #
| Absent(*The image has no function of that name.
*) - #
| 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 = The type for a driver's answer when asked to grow a device's queues (Driver.grow).
Drivers
Opening
type 'a opener = {- #
name : string;(*The device's name, without a machine suffix.
*) - #
driver : (module Driver with type t = 'a);(*The device's driver.
*) - #
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).