Errors, memory and threads

In short: recoverable failures return Status or Expected<T>, and bugs abort through RML_CHECK; nothing throws. Tensor buffers come from AlignedAlloc, which hands out 64-byte-aligned blocks and keeps four process-wide counters. Arena is a bump allocator over one such block, and a thread-scoped BufferPool recycles freed tensor buffers by size. ThreadPool::ParallelFor cuts a range into chunks whose boundaries depend only on the range, the grain and the thread count. Threads claim chunks from an atomic counter, idle workers spin briefly before sleeping, a dispatch allocates nothing, and a nested call runs inline.

The design notes are the Errors, Memory and Threads sections of docs/architecture.md. RML_CHECK itself is covered on Tensors and dtypes. For background, see GPU execution (which thread_pool.h links to) and Reproducibility in the ML models reference.

The error model

rml splits failures in two. A failure caused by something outside the caller's control, such as a missing file, a malformed header, an unsupported config or a failed large allocation, is returned as a value. A failure caused by the caller breaking a contract, such as a shape mismatch, is a bug and aborts. The header comment states the rule (L4–6): hot paths never throw.

Status

Expected<T>

std::expected is C++23 and rml targets C++20, so Expected<T> is the subset rml needs, built on std::variant<T, Status>. Index 0 is the value and index 1 the error.

Propagation macros

Without exceptions, each fallible call needs a check and an early return. Two macros write that out.

Loading a model is the heaviest user. DecoderLM::FromHf runs, in order, RML_ASSIGN_OR_RETURN(ModelConfig config, LoadHfConfig(...)) and RML_ASSIGN_OR_RETURN(io::SafetensorsModel model, io::SafetensorsModel::Open(directory)). Either failure returns to the caller unchanged, with its code and message. See The decoder and Model files.

Aligned allocation

Every tensor buffer is 64-byte aligned: one cache line on x86 and Apple Silicon, and enough for any SIMD load (L6–9). Storage::Allocate, which backs every Tensor rml allocates itself, calls AlignedAlloc, unless a BufferPool is installed on the thread.

AlignedAlloc over-allocates with malloc and rounds up inside the block. Just below the pointer it returns, it stores a two-word Header holding the original malloc pointer and the requested size. AlignedFree needs nothing else, so there is no lookup table and no lock. Before calling malloc, it checks that bytes plus the overhead can't overflow (L40–41). HugeRequestFails asks for SIZE_MAX − 8 bytes and expects nullptr.

The counters

AllocationStats is four process-wide numbers: successful allocations, frees, live bytes and peak live bytes. They are relaxed atomics and always on. Only the order of each counter's own updates matters, so no stronger memory order is needed. The peak is raised with a compare-and-swap loop that retries only while the new live total is still above the stored peak. ResetPeakAllocation sets the peak back to the current live bytes, so a test can measure the peak of one region.

These counters only see AlignedAlloc. Memory that is not tensor storage, like a std::vector or mmapped weights wrapped with Tensor::FromBuffer, is not counted. So the decoder's zero-allocation test uses a different tool: DecodingAllocatesNothing replaces the global operator new with a counting one. It loads a model, runs a prompt and one warm-up step, and then expects 16 more decode steps to add no allocations, with and without a thread pool. See Models and memory for the memory budget this protects.

Arena

Arena is a bump allocator over one block from AlignedAlloc. Allocating moves an offset forward, and memory comes back all at once.

AllocatingFromArenaDoesNotTouchHeap resets and allocates 100 times and checks that the allocation counter does not move. The design notes say per-step scratch comes from an arena reset at each step (L66–67). At this commit, though, no code outside rml/core uses Arena, so it is available rather than wired into decoding. The graph executor reuses memory through a BufferPool instead.

BufferPool

A model run with static shapes allocates the same sizes every time. BufferPool caches freed tensor buffers keyed by exact size, so later runs reuse warm memory instead of mapping and faulting in fresh pages for every op (L4–11). It is opt-in per thread: while a ScopedBufferPool is alive, every Storage::Allocate on that thread draws from the pool.

The hand-back happens in tensor.cc. A pooled Storage::Allocate wraps the block with a PooledBlock owner that holds a shared_ptr to the pool and calls Release in its destructor. So a tensor that outlives the scope still returns its buffer when it dies, and it keeps the pool alive until then. Blocks come from AlignedAlloc, so a cache miss shows up in the allocation counters and a reused block does not. The one caller at this commit is the graph executor, which keeps a shared_ptr<BufferPool> and opens a ScopedBufferPool at the start of each run (executor.cc L156); see Executing a graph.

ThreadPool

ThreadPool starts n − 1 worker threads that spin for a short while after each loop and then sleep on a condition variable. The thread calling ParallelFor works too, so ThreadPool(1) has no workers and runs everything inline. ThreadPool::Default() is a process-wide pool sized by RML_NUM_THREADS, or else by hardware_concurrency(). It is created on first use and intentionally never destroyed.

How a range is chunked

1. NumChunks returns 1 when the pool has no workers or the calling thread is already running a chunk. Otherwise it returns n / grain, clamped to between 1 and 4 × num_threads.
2. With one chunk, ParallelFor calls fn(begin, end) directly and never touches the pool (L54–57).
3. Otherwise chunk_size = ⌈n / num_chunks⌉, and the job runs ⌈n / chunk_size⌉ chunks. Chunk i covers [begin + i·chunk_size, min(that + chunk_size, end)). Only the last chunk can be short. Since num_chunks ≤ n / grain, every other chunk has at least grain elements.

For example, ParallelFor(0, 1000, 10, fn) on a 4-thread pool gives 1000 / 10 = 100, clamped to 16 chunks. chunk_size is ⌈1000 / 16⌉ = 63, so chunks 0 to 14 are 63 elements and chunk 15 is [945, 1000), 55 elements. With (0, 100, 1), the same pool computes 16 chunks of ⌈100 / 16⌉ = 7, but 7 × 15 is already 105, so it runs 15 chunks and the last is [98, 100).

What is deterministic

Chunk boundaries depend only on begin, end, grain and the thread count (L8–11). A kernel that writes each output element from inside one chunk, and reduces in a fixed order within it, gets bit-identical results however the chunks are scheduled. Changing the thread count changes the boundaries, so a reduction split across chunks can round differently on 4 threads than on 8. RML_NUM_THREADS pins it. See Reproducibility for why floating-point order matters.

Dispatch without allocation

ParallelFor takes any callable, but the pool can't store a template type. Instead it builds a Job on the caller's stack: a plain function pointer that casts a void* back to the callable, plus the chunk layout. No std::function, so no heap allocation. Run then:

1. Locks call_mutex_, so two threads calling ParallelFor on the same pool take turns rather than sharing workers.
2. Publishes the job, resets next_chunk_ to 0 and sets pending_workers_ to the worker count, all atomics, then bumps generation_. Spinning workers see the new generation on their own; only if sleeping_workers_ is above zero does it take sleep_mutex_ and notify the condition variable.
3. Works on chunks itself through WorkOn, which claims the next index with fetch_add until none are left.
4. Spins until pending_workers_ reaches zero, yielding its core after kSpinsBeforeYield = 256 polls (L23–25). It waits for every worker to leave the job, even ones that found no chunk left: the job lives in this stack frame, so returning earlier would leave a worker holding a dangling pointer.

A worker first spins on generation_ for up to RML_THREAD_SPIN_US, checking the clock every 64 polls, since in a model the next parallel loop usually follows within microseconds. If nothing arrives it increments sleeping_workers_ and waits on the condition variable until generation_ changes or stop_ is set. It then runs WorkOn and decrements pending_workers_ with release order, which the caller's acquire load pairs with. Each worker remembers the last generation it ran, so a wakeup without a new job sends it back to waiting. The generation_ bump and the sleeper's increment are sequentially consistent, so either the sleeper sees the new job or the caller sees it asleep and wakes it (L79–88).

Nested calls

A thread_local flag, in_parallel_region, is set while a thread runs chunks. A ParallelFor called from inside a chunk sees it, gets one chunk from NumChunks, and runs inline. So a parallel kernel can call another parallel kernel, for a kernel that calls GEMM from inside its own chunks, without deadlocking on call_mutex_ or oversubscribing the cores. The outer loop gets the parallelism.

Who uses it

Tests

How these run under sanitizers is on Testing.