rig

rig drives Apple GPUs through Metal, NVIDIA GPUs through CUDA or NVIDIA's kernel driver, and AMD GPUs through Linux's amdgpu. On Linux it also drives AMD and NVIDIA GPUs with no kernel driver, over PCI. Files are devices of their own, and another machine's devices open through an agent running there.

Install

rig needs OCaml 5.5 or later. Install the rig package and add rig to your executable's libraries:

$ opam install rig

(executable
 (name main)
 (libraries rig))

Each driver is a library of its own, such as rig.metal, rig.cuda or rig.amd, which a program adds beside rig. Some open their devices through a further library, such as rig.amd.amdgpu for AMD GPUs under Linux's amdgpu.

A first program

submit returns the point of the work on its device's timeline, and Point.wait returns once the work up to that point has completed. The preamble of rig.mli defines timelines, values and points.

This program copies one integer on a memory device, whose memory is the host's and whose work runs before submit returns. It runs on any machine.

# open Rig;;

# let () =
    let device = Result.get_ok (memory_device "example") in
    let src = Buffer.create device 4 in
    let dst = Buffer.create device 4 in
    Buffer.write src Bigarray.int32 (fun a -> a.{0} <- 42l);
    let work =
      Submission.(make device [| copy ~src:(Buffer src) ~dst:(Buffer dst) |])
    in
    let run = Submission.take work in
    let done_ = submit run ~buffers:[||] in
    Point.wait done_;
    Buffer.read dst Bigarray.int32 (fun a ->
        Printf.printf "copied %ld\n" a.{0});;

copied 42

Each call, in order:

• memory_device "example" opens the memory device named "example", or returns it if it is open already.
• Buffer.create device 4 allocates 4 bytes of the device's memory.
• Buffer.write src Bigarray.int32 f calls f on src's bytes as an array the host writes, since the host addresses a memory device's memory, once the work that used them is done.
• copy ~src:(Buffer src) ~dst:(Buffer dst) is one part of the work, a copy between two buffers that every submit uses. An end may instead be Slot i, the buffer buffers.(i) of each submit (Submission.side). rig orders the parts of a submission by the memory they use. It runs a copy on the queue of the last earlier part the copy follows, where that queue runs copies, and otherwise on the device's copy engine, the queue its driver names in the copy_engine fact (Submission.make).
• Submission.make device [| ... |] prepares the parts once, as a submission that any number of submits run.
• Submission.take work takes a run of the submission: the storage that a submit's work reads its parameters from, such as a launch's grid (On a GPU). A copy has none, but every submit takes a run, and the submit frees it for the next.
• submit run ~buffers:[||] runs the run's submission once. buffers holds the buffers that change from one submit to the next, which a launch's parameters and a copy's Slot ends name by index. This copy names none. An optional ~waits lists points of other devices the work must follow.
• Point.wait done_ returns once the copy has completed. Buffer.read would wait for it too, since it follows the last write of dst.

On a GPU

On a GPU, the work is usually a launch, which runs a function of an image, a binary compiled for the device, over a grid of threads. Image.resolve finds the function and makes it ready to launch. The calls are the same on every GPU, and the work may complete after submit returns.

This function doubles the n floats of the device buffer x, n a multiple of 256, with the function twice of binary:

let double d binary x n =
  let image = Rig.Image.load d binary |> Result.get_ok in
  let twice = Rig.Image.resolve image "twice" |> Option.get in
  let twice =
    Rig.Submission.(launch ~refs:[| write 0 ~at:0 |] twice ~params:8)
  in
  let s = Rig.Submission.make d [| twice |] in
  let run = Rig.Submission.take s in
  let b = Rig.Submission.block s 0 in
  Rig.Submission.Run.groups run b (n / 256) 1 1;
  Rig.Submission.Run.threads run b 256 1 1;
  Rig.Point.wait (Rig.submit run ~buffers:[| x |])

The launch's 8 bytes of parameters hold one pointer. A ref (Submission.ref) makes 8 bytes of the parameters an offset into one of the submit's buffers, its slot. This one, write 0 ~at:0, makes bytes 0 to 7 an offset into buffers.(0), here x, which the work reads and writes, and the driver's submit adds x's address to it as it encodes the work (Drivers). The run never stores the offset, and bytes a run never stores read 0, so twice gets x's address.

A ref states how the work uses its slot: Submission.read, Submission.write or Submission.unordered. Submission.make derives the slots from the refs and the copies' Slot ends, and each submit orders the work after earlier work on each slot's memory by that use.

A block is the part of a run that holds one launch's grid, the threads of each group, and its parameters (Submission.block). Submission.block s 0 names the block of part 0, and the two stores set its grid to n / 256 groups of 256 threads.

On a Mac, the binary is a metallib, compiled with xcrun -sdk macosx metal -c twice.metal -o twice.air and xcrun -sdk macosx metallib twice.air -o twice.metallib from:

struct args { device float *x; };
kernel void twice(constant args &a [[buffer(0)]],
                  uint i [[thread_position_in_grid]]) { a.x[i] *= 2; }

With CUDA, it can be PTX, which CUDA compiles for the GPU as it loads, made with nvcc -ptx -o twice.ptx twice.cu from:

extern "C" __global__ void twice(float *x) {
  x[blockIdx.x * blockDim.x + threadIdx.x] *= 2;
}

A program opens the GPU through its driver's opener, moves the floats to the GPU's memory, runs double, and copies them back:

let () =
  let d = Rig.open_ (Rig_metal.gpu 0) |> Result.get_ok in
  (* With CUDA: Rig.open_ (Rig_cuda.gpu 0) and "twice.ptx". *)
  let binary =
    In_channel.with_open_bin "twice.metallib" In_channel.input_all
  in
  let n = 65536 in
  let host = Rig.Buffer.create Rig.host (4 * n) in
  Rig.Buffer.write host Bigarray.float32 (fun a ->
      for i = 0 to n - 1 do a.{i} <- float_of_int i done);
  let x = Rig.Buffer.create d (4 * n) in
  Rig.Buffer.copy ~src:host ~dst:x;
  double d binary x n;
  Rig.Buffer.copy ~src:x ~dst:host;
  Rig.Buffer.read host Bigarray.float32 (fun a ->
      Printf.printf "%g\n" a.{n - 1})

Buffer.copy orders itself after the work that used its buffers, so the copy back follows the launch (the Order section of rig.mli, and Ordering). A copy returns once the bytes are in its destination, and Buffer.read follows the last write of its buffer whatever came before it (the Buffer module's Host access section). Memory traces each allocation's stamps through this program, and through a loss of the GPU.

Reading order

The interfaces state every rule, and read top to bottom:

1. Rig: devices, timelines, the order of work, buffers, running work, claims, domains, loss and the return of memory.
2. Rig_edge: the contract a driver implements.
3. One driver: Metal, CUDA, NVIDIA, or AMD through amdgpu.

These pages show the rules at work, with programs and traces:

1. Ordering: how work on several devices is ordered by the memory it uses.
2. Claims: a function that writes over its argument.
3. Memory: a buffer's life, and its stamps through the GPU program above.
4. Loss: what a program sees when a device fails.
5. Drivers: one submission from submit to the hardware and back.
6. GPUs without a driver: opening an AMD or NVIDIA GPU over PCI.
7. Machines: devices of other machines.

The numbered examples, from 01-add, teach one idea each and run on any machine. x-pci needs Linux. Hardware notes describe the hardware behaviour each driver relies on.

Examples

• 01-add ↗A kernel is a function that a device runs once for each thread of a grid.
• 02-buffers ↗A buffer is bytes of one device's memory, and the caller gives them their meaning.
• 03-files ↗The disk is a device whose buffers are files.
• 04-timelines ↗A device numbers the work submitted to it 1, 2, 3, ..., the values of its timeline.
• 05-stamps ↗Each allocation records the point of its last write and, for each device, the point of its last use.
• 06-steps ↗Compiled work is made once and run many times.
• 07-reclamation ↗Nothing frees a buffer by hand.
• 08-claims ↗A function handed buffers may write over an input in place only if nothing else reads that memory.
• 09-loss ↗A device whose work fails is lost, once and for good.
• 10-profiles ↗A profile holds what every device did while a function ran.
• 11-elf ↗Host programs, GPU programs and firmware come as ELF objects.
• 12-host-programs ↗Host code is a C function compiled to an ELF object, loaded as an image on the host and launched as a GPU's kernel is.
• 13-pool ↗One pool of threads per process computes on the host, shared by nx's kernels, host programs and a library's own C code.
• 14-machines ↗A job is one process, the controller, and an agent on each other machine, which opens that machine's devices and runs the work the controller submits.
• x-pci ↗Needs Linux.

this page is tested with mdxedit this page ↗