Module Rig_metal

Each device of this library owns a Metal command queue on the Mac's GPU. A device has one rig queue, "COMPUTE:0", which runs launches and fills. Its memory is shared memory that the GPU and the host address alike, so the host copies it and the device has no copy engine.

A kernel reaches the GPU in a metallib, the binary Apple's Metal compiler makes. Rig.Image.load loads the metallib as a Metal library. The first Rig.Image.resolve of one of its functions compiles the function for the GPU into a pipeline, which later launches reuse, unless Rig.Image.prepare started the compile before; Metal compiles the functions prepared together at once. A Metal command buffer holds encoded work, and Metal runs it only once it is committed. The device keeps at most one open command buffer, which collects the work of consecutive values until the device commits it. Each Rig.submit encodes its launches into the open command buffer and returns without waiting for them. The device commits the open command buffer when a wait asks for one of its values, and on its own often enough to keep the GPU busy (Work). As each command buffer completes, Metal calls a handler of the device on one of Metal's threads. The handlers write into the device's timeline word the last value of each command buffer once it and every buffer committed before it have completed without failure.

This program adds x to y on the GPU and prints y[10] = 11. It needs add.metallib, compiled from

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

let n = 1024

let () =
  let ok = function Ok v -> v | Error why -> failwith why in
  let d = ok (Rig.open_ (Rig_metal.gpu 0)) in
  let lib =
    In_channel.with_open_bin "add.metallib" In_channel.input_all
  in
  let image = ok (Rig.Image.load d lib) in
  let x = Rig.Buffer.create d (4 * n)
  and y = Rig.Buffer.create d (4 * n) in
  Rig.Buffer.write x Bigarray.float32 (fun a ->
      for i = 0 to n - 1 do a.{i} <- Float.of_int i done);
  Rig.Buffer.write y Bigarray.float32 (fun a -> Bigarray.Array1.fill a 1.);
  let add = Option.get (Rig.Image.resolve image "add") in
  let add =
    Rig.Submission.(
      launch ~refs:[| read 0 ~at:0; write 1 ~at:8 |] add ~params:16)
  in
  let s = Rig.Submission.make d [| add |] 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;
  ignore (Rig.submit run ~buffers:[| x; y |] : Rig.Point.t);
  Rig.Buffer.read y Bigarray.float32 (fun a ->
      Printf.printf "y[10] = %g\n" a.{10});
  Rig.close d

The host addresses the device's memory, so it writes x and y in place, after any work on them (Rig.Buffer.write). The program never stores the 8 parameter bytes at 0 and 8, so they hold the offset 0, and the kernel reads them as the GPU addresses of x's and y's first bytes, a.x and a.y. n is a multiple of 256, the threads of a group. Rig.Buffer.read of y returns once the launch has completed, and raises Rig.Lost if its command buffer failed.

A device lives until Rig.close or a loss (Failure and stopping). Any domain may call any function of this module, at the same time as other domains. sleep and the first entry for a function let other domains run while they wait for Metal.

An icb (Rig_metal_abi) and an opener's make, called from a callback inside a lock (Rig's domains section), raise Invalid_argument, having changed nothing. rig calls the functions of The driver, but locate, peer and signaled, holding no lock (Drivers); one called from such a callback raises Invalid_argument, having changed nothing, where it would take a lock that the interrupted call holds.

Compiled code can encode Metal commands of its own, such as indirect command buffers, in a fill: a C function that the driver's submit calls with the open command buffer's encoder (Rig_metal_abi). Rig_edge states how rig drives a device through its driver, and The C functions how it hands the driver work.

Opening

A device opens on macOS 15 or later, on a GPU of an Apple or a Mac GPU family. The library builds on every system, and off macOS no device opens.

The driver

rig calls these functions. On Metal, alloc, free, locate, map_host, peer, map_peer and dma behave as in Memory; image, prepare, entry and unload as in Code; sleep and stop as in Failure and stopping; and the C functions as in Work.

Facts

arch is the highest Apple GPU family the GPU supports, such as "Apple7" for an M1, else its Mac family, "Mac2". budget is the memory, in bytes, that Metal recommends the device keep allocated at most. The word holds an unsigned 64-bit integer in the host's byte order, and the device's values start after the value it holds when make returns. Its location has a host address and no device address: only the completion handlers and stop write it.

The device's queue waits only for the device's own work, so rig waits on the host for the points a submission waits for, before it hands the submission to the driver. hang_ms is None because Metal ends a command buffer it judges to run too long, with MTLCommandBufferErrorTimeout.

Work

The device's struct rig_driver holds rig_metal_room, rig_metal_submit and rig_metal_commit, which rig_metal.h declares, and no settle: the command buffers' completion handlers write a value's times before they raise the word. rig_metal_submit encodes the parts of a value, in order, into the open command buffer, and opens one if none is open. A command buffer has one compute encoder, which runs its commands one after the other (MTLDispatchTypeSerial), and each command buffer runs after the one committed before it.

The device commits the open command buffer in four cases:

• a commit (The C functions), for any value it is given;
• a rig_metal_submit that finds fewer than three of the device's committed command buffers uncompleted, so that the GPU finds work queued behind the one it runs;
• the completion of a command buffer while the open one holds work;
• a rig_metal_submit that leaves 256 values in it, so the device commits each value at most 256 values late.

A rig_metal_submit that commits returns RIG_COMMITTED.

Metal calls a handler of the device on one of Metal's threads once a command buffer has completed, successfully or not (addCompletedHandler:). The handlers store into the word in the order the command buffers were committed, with release ordering. Once a command buffer and every one committed before it have completed, the last value whose work it holds is stored. A command buffer that failed stores nothing, and from then on no handler stores into the word, even for a later command buffer that succeeds, until stop. A value with no parts is reached once the work before it has completed.

A fill receives the open compute encoder, encodes into it, and may start new command buffers (Rig_metal_abi).

rig_metal_room returns RIG_LATER and RIG_GROW for launches' parameters alone (Launches). For command buffers, rig_metal_submit instead blocks, before it opens a command buffer, while the device has 1,024 command buffers, as many as its Metal queue holds, counting the open one and the committed ones not yet completed. rig_metal_commit commits the open command buffer, which holds the work of every value not yet committed.

Launches

A launch dispatches its function once, in the open command buffer's encoder. It sets the function's pipeline; its parameters as kernel buffer 0, unless it has none, in the device's parameter arena, a shared buffer bound at each launch with setBuffer:offset:atIndex:; its threadgroup memory, rounded up to a multiple of 16 bytes, as threadgroup memory 0 (setThreadgroupMemoryLength:atIndex:), unless it is 0; and its grid (dispatchThreadgroups:threadsPerThreadgroup:). The function reads its parameters' bytes at [[buffer(0)]], in the constant or the device address space, and each ref's 8 bytes as the GPU address of the ref's buffer plus the offset. Its parameters start on 256 bytes, so any type of an alignment of at most 256 bytes may be declared there. The Metal compiler assumes that the object a pointer argument refers to is aligned as its type requires (Metal Shading Language Specification, §2.5), and Apple states no alignment for the copy setBytes:length:atIndex: makes, so a launch never binds its parameters with it.

The Metal Feature Set Tables bound a buffer's contents by its length alone, so the device's maxBufferLength bounds a launch's parameters. The launches take the arena's bytes in turn, each at an offset that is a multiple of 256 bytes, the largest constant buffer offset alignment, from an arena whose first byte is on 256 bytes. The bytes of a command buffer's launches come back once its work completes. rig_metal_room returns RIG_LATER for parts whose parameters the arena holds only once earlier work completes, and RIG_GROW for parts whose parameters exceed the empty arena, as a device's first launch with parameters finds it. The growth (grow) makes an arena of the least power of two of bytes, at least 1 MiB and at most maxBufferLength, that holds twice them, and the next submission takes it; the command buffers that use the old one keep it until their work completes. An arena of 1 MiB holds the parameters of 4,096 launches of at most 256 bytes each whose command buffers have not completed. A submission past what the arena holds waits until enough of that work completes, and the GPU runs it meanwhile. A growth Metal has no memory for answers Short, so a launch rig cannot find memory for raises Rig.Out_of_memory. Apple documents no alignment for a buffer's first byte (Memory), so a growth whose arena Metal places, by its GPU or host address, off a multiple of 256 bytes raises Fault and makes nothing. The arena never shrinks: it stays until the device stops, and counts in no budget, since a budget counts buffers, images' code and caches (Reclamation), as no driver's queues count in one.

rig_metal_room refuses a launch if:

• an axis of its grid or threadgroup has size 0;
• its threadgroup holds more threads, along an axis or in all, than the least of the pipeline's maxTotalThreadsPerThreadgroup and the three axes of the device's maxThreadsPerThreadgroup;
• its threadgroup memory, rounded, exceeds the device's maxThreadgroupMemoryLength less the function's staticThreadgroupMemoryLength.

It also refuses parts whose parameters take more than maxBufferLength together.

Profiles

While a profile is taken, rig_metal_submit commits the open command buffer, runs the value's work in command buffers of its own and commits them immediately. The value's times are the first one's GPUStartTime and the last one's GPUEndTime. Profiling therefore costs the device its batching.

Memory

alloc makes the same shared memory (MTLResourceStorageModeShared) for every Rig_edge.memory. A region is resident for the work of every submission whose submit starts after alloc or map_host has returned. Apple documents no alignment for a buffer's first byte. Metal places a buffer below a page at a multiple of 256 bytes and a larger one at a page, so the device's alignment (Rig_edge.facts) is 256. A buffer whose GPU or host address breaks it faults the device, so Rig.Buffer.create raises Rig.Lost, and so does every later use of the device and of its memory. alloc and map_host raise Invalid_argument for fewer than one byte.

locate returns a region's GPU address (gpuAddress), which kernels read and write, its host address, and its MTLBuffer. free takes the region out of the device's residency set, unless the device stopped, and releases its MTLBuffer.

Metal wraps only memory that starts at a page, so map_host d p n wraps the pages that hold the n bytes at p, without a copy, and the region's host address is p rounded down to a page. It is Never if Metal cannot wrap the memory; Metal gives no reason.

peer is false and map_peer is Never, since a Metal device maps only its own memory and the host's. dma is Error, since Metal gives no other device DMA to its buffers.

Code

image loads a metallib as a Metal library, and leaves compiling its functions to prepare and entry. The result is Ok (Loaded i), or Error with the reason if Metal cannot load the bytes as a library, or if one of its functions is not a compute kernel.

prepare starts Metal's asynchronous compile of a function into a pipeline (MTLComputePipelineState) that indirect command buffers may also run, and returns. Metal compiles in its own processes, several pipelines at once. The first entry for a function waits for the compile prepare started, or compiles the same way and waits. A prepare or an entry waits while another domain's entry waits for a pipeline of the same image. Calls for one function from several domains concurrently make one pipeline. Every later call returns the same entry, which is valid until unload. The entry's code is the pipeline's address, and its launch the address of what a launch reads, the pipeline and its limits. Its threads are the least of the pipeline's maxTotalThreadsPerThreadgroup and the device's axes. Its params are the size, trailing padding included, of the struct the function reads as kernel buffer 0, Some 0 for a function with no buffer 0, and None where buffer 0 points to another type, such as words, whose length the function does not state. entry raises Invalid_argument with Metal's reason if Metal cannot make a pipeline of the function, such as for a function that needs more threadgroup memory than the GPU family has. A refusal is not kept, so a later call compiles again.

unload releases the library, and the pipelines and launches entry made of it. It waits for no compile: a pipeline prepare started and no entry took is released once made. An indirect command buffer that records one of these pipelines keeps it until its own release (Rig_metal_abi.icb).

Failure and stopping

A submission fails immediately if a fill returns a failure, if Metal cannot make a command buffer or an encoder, or if Metal raises an exception. rig_metal_submit then drops the open command buffer uncommitted, with the work of every value it held, and returns RIG_FAILED with the reason. It fails later if Metal reports that a command buffer holding its work failed, such as one Metal ended because it ran too long (MTLCommandBufferErrorTimeout) or kept the GPU from the display. The reason is then "the GPU's work failed: " followed by Metal's description of the error. Either way no handler stores into the word again (Work), and the word stays below every value the failed command buffer held until stop.

The device keeps its first failure. Once it has recorded the failure, sleep raises Fault with it, every rig_metal_submit returns RIG_FAILED with it and runs nothing, and every rig_metal_commit returns RIG_FAILED. rig loses the device at the first such call, which these make:

• a Rig.submit on the device;
• a wait for the device's work, such as Rig.Point.wait, Rig.Buffer.wait or a Rig.close that has work to wait for.

Until then the device is live. Work encoded before the failure was recorded may still run. While a command buffer runs, a wait lasts until the buffer completes or Metal reports its failure.

sleep blocks on a condition the handlers signal, and uses no processor time while it waits.

stop ignores fault. It drops the open command buffer uncommitted, so its work never runs. If every command buffer the device committed has completed, it writes the last value rig_metal_submit received into the word, and no work of the device writes memory again. Otherwise the committed work may write memory while it runs. Once the last such command buffer has completed, its handler writes the last value into the word, whatever that work did. The stop gives back what the device runs work with: its Metal command queue, fence, residency set and command buffer ring, which each command buffer, committed or dropped, keeps until Metal releases its completion handler, some time after the handler returns. Once the word shows the last value and stop has returned, no handler writes it: the handlers write in order, and the one that completes the last command buffer writes the last value. The few bytes of C state that read the word stay with the device. sleep on a stopped device raises Invalid_argument.

References

• Apple's Metal framework headers (macOS 26 SDK): MTLCommandQueue.h (newCommandQueueWithMaxCommandBufferCount:, addResidencySet:), MTLCommandBuffer.h (addCompletedHandler:, status, error, MTLCommandBufferErrorTimeout, GPUStartTime, GPUEndTime, computeCommandEncoderWithDispatchType:), MTLFence.h, MTLResidencySet.h (addAllocation:, commit), MTLDevice.h (newBufferWithLength:options:, newBufferWithBytesNoCopy:length:options:deallocator:, recommendedMaxWorkingSetSize, maxBufferLength, supportsFamily:, maxThreadsPerThreadgroup, maxThreadgroupMemoryLength), MTLBuffer.h (gpuAddress, contents), MTLLibrary.h (newLibraryWithData:error:, functionNames, newFunctionWithName:), MTLFunction.h (functionType), MTLComputePipeline.h (supportIndirectCommandBuffers, maxTotalThreadsPerThreadgroup, staticThreadgroupMemoryLength), MTLDevice.h (newComputePipelineStateWithDescriptor:options:completionHandler:), MTLComputeCommandEncoder.h (setBuffer:offset:atIndex:, setThreadgroupMemoryLength:atIndex:, dispatchThreadgroups:threadsPerThreadgroup:), MTLIndirectCommandBuffer.h.
• Metal Best Practices Guide, Pipelines: "Building these pipelines asynchronously maximizes performance and parallelism."
• Simplifying GPU resource management with residency sets: a set added to a queue keeps its committed allocations resident for every command buffer of the queue.
• makeBuffer(bytesNoCopy:length:options:deallocator:): the wrapped memory starts at a page.
• setBytes(_:length:index:): it states no alignment for its copy.
• setBuffer(_:offset:index:): the offset of a buffer bound to a device argument is aligned to its type's size, and a constant one's to the family's minimum.
• Metal Shading Language Specification, §2.5: the compiler assumes that the object a pointer argument refers to is aligned as its type requires.
• Metal feature set tables: GPU families, minimum constant buffer offset alignment, and buffer contents limited only by the buffer's length.