Moirai Parallel Execution Backend

Moirai is the Atlas-native hybrid concurrency library that provides the parallel execution backend for batched voxel-wise fitting in the diffusion MRI pipeline. It blends synchronous data-parallel operators (the Rayon-replacement surface) with a lock-free work-stealing scheduler, a unified hybrid runtime that serves both async and parallel work, and compile-time execution policies that monomorphize to zero-cost concrete paths.

At the RITK level, Moirai is surfaced through the MoiraiBackend type in ritk-image and the broader Coeus backend traits. Compute-heavy diffusion kernels — DTI log-linear solves, DKI Levenberg-Marquardt iterations, CSD non-negative deconvolution — distribute across CPU cores through Moirai’s data-parallel operators. The important design point is that Moirai does not require a separate image type or parallel-only API: readers, writers, filters, and registration code are parameterized over a backend, so the same Image<f32, B, 3> contract can execute on the sequential backend for deterministic single-threaded work or on Moirai for throughput.

Architecture

Moirai has three layers, each a separate domain crate:

A .par() worker task runs on the same unified hybrid scheduler as async work, not a separate thread pool. The sync return shape is a property of the operation, not an isolation boundary: a parallel kernel can spawn or drive async work on the same runtime, and sync operators compose cleanly inside pure compute kernels without async contagion.

The scheduler uses lock-free work-stealing: each worker owns a private Chase-Lev deque and a SplitDeque that keeps the freshest work on a local stack (zero atomics) and spills only the oldest half to a shared deque for thieves. Batch stealing transfers up to 16 items at a time, reducing steal contention when spawn rate greatly exceeds steal rate.

Execution Policies

Every data-parallel operator is generic over an ExecutionPolicy — a zero-sized type-level marker that decides whether to parallelize. Because the decision is an associated function, generic code over P: ExecutionPolicy monomorphizes to one concrete path with no value passed and no dynamic dispatch.

The Adaptive policy is the everyday entry point. It parallelizes only at or above 1,024 elements, where dispatch and join overhead is amortized; below that threshold the sequential path eliminates all scheduling cost. Because Adaptive is itself a zero-sized type, .par() is a fully monomorphized, zero-cost abstraction — the parallel/sequential decision is a cheap inlined runtime check.

use moirai_parallel::{Adaptive, Parallel, Sequential, ExecutionPolicy};

// Force sequential (deterministic / nested)
for_each_with::<Sequential, _, _>(&data, |x| process(x));

// Force parallel (ignore threshold)
for_each_with::<Parallel, _, _>(&data, |x| process(x));

// Adaptive default (everyday use)
data.par().for_each(|x| process(x));

Extension Traits

The primary surface is two extension traits on slices:

ParallelSlice<T>

pub trait ParallelSlice<T> {
    /// Adaptive, auto-routing parallel view over `&[T]`.
    fn par(&self) -> ParRef<'_, T, Adaptive>;
}

ParRef provides read-only parallel combinators:

ParallelSliceMut<T>

pub trait ParallelSliceMut<T> {
    /// Adaptive, auto-routing mutable parallel view over `&mut [T]`.
    fn par_mut(&mut self) -> ParMut<'_, T, Adaptive>;
}

ParMut provides mutable in-place combinators:

use moirai_parallel::{ParallelSlice, ParallelSliceMut};

let v: Vec<u64> = (0..1_000_000).collect();

// Map-reduce over a read-only slice
let sum = v.par().map_reduce(0u64, |&x| x, |a, b| a + b);

// In-place mutation of every element
let mut doubled = v.clone();
doubled.par_mut().for_each(|x| *x *= 2);

Data-Parallel Free Functions

For rare cases that need to force a specific policy via turbofish, every operator is also available as a *_with::<P> free function. These are the low-level override and are what the extension traits delegate to.

Read-only operators

Mutable operators

Index-domain operators

use moirai_parallel::{for_each_with, map_collect_with, map_reduce_with, Sequential, Parallel};

// Force sequential: deterministic, no scheduling overhead
for_each_with::<Sequential, _, _>(&small_data, |x| debug_assert!(x > 0));

// Force parallel: even for small inputs (benchmarking, nested regions)
let results: Vec<_> = map_collect_with::<Parallel, _, _, _>(&data, expensive);

// Elementwise product across multiple slices (index-aligned)
let dot = reduce_index_with::<Adaptive>(
    n, 0.0_f64,
    |i| a[i] * b[i],
    |x, y| x + y,
);

Chunk Operators

For batched workloads — like per-voxel model fitting where each “chunk” is a contiguous slice of the voxel array — chunk operators amortize dispatch overhead by assigning each worker a range of consecutive elements. The diffusion pipeline typically uses for_each_chunk_mut_with_state to give each worker a reusable scratch buffer (design matrix, workspace arrays) across its assigned voxels.

use moirai_parallel::for_each_chunk_mut_with_state;

// Per-voxel DTI fitting: each worker gets a private design matrix
for_each_chunk_mut_with_state::<Adaptive, _, _, _, _>(
    &mut voxel_results,
    256,                             // chunk size (voxels per task)
    || VoxelWorkspace::new(),         // init per-worker scratch
    |ws, chunk| {
        for result in chunk {
            *result = fit_dti(ws, &signal[result.index]);
        }
    },
);

The multi-buffer variants (pair, triple, quad) are fused kernels that update several output arrays in one authoritative pass — for example, writing FA, MD, and PEV fields simultaneously — without allocating intermediate tuples.

Borrowing Scope

scope provides a Rayon-style region for spawning parallel sub-tasks that capture non-'static references. It blocks until every spawned job has completed, so borrowed data cannot escape.

use moirai_parallel::scope;

let data: Vec<u64> = (0..1000).collect();
let sum = std::sync::atomic::AtomicU64::new(0);

scope(|s| {
    s.spawn(|| sum.fetch_add(data.iter().sum::<u64>(), Ordering::Relaxed));
    s.spawn(|| sum.fetch_add(data.len() as u64, Ordering::Relaxed));
});
// Both tasks guaranteed complete; `data` is still alive.

Join

join runs two closures concurrently. The left closure is scheduled on a worker thread; the right closure runs on the caller. If the scheduler refuses the left closure (shutting down or admission queue full), it runs on the caller instead, so join gracefully degrades to sequential but never drops work.

use moirai_parallel::join;

let (a, b) = join(
    || expensive_first_half(),
    || expensive_second_half(),
);

For explicit policy selection, use join_with::<P>.

Work-Stealing Deques

The scheduler’s lock-free work-stealing deques live in moirai-scheduler:

The split-deque design is the key performance optimisation for the diffusion pipeline: an owner’s push/pop against its private stack costs zero atomics, and only work old enough to be worth stealing reaches the shared deque. A thief that finds the victim empty returns StealResult::Empty (look elsewhere) rather than spinning; a thief that loses a race returns StealResult::Retry (the victim may still hold work, so retrying the same deque is worthwhile).

Batch stealing (steal_batch) transfers half the available items (up to 16) in one atomic CAS, cutting steal contention by an order of magnitude when many small tasks are spawned rapidly — the diffusion pipeline’s voxel-chunk dispatch pattern.

Memory ordering follows the weak-memory-correct Chase-Lev formulation of Lê, Pop, Cohen & Nardelli (PPoPP 2013), with a Morrison-Afek fence-free fast path for pop on x86/x86_64 when the deque holds enough items that no steal can reach the popped slot. Old buffers freed by resize are retired to an epoch-reclamation list and freed only when no accessor is in-flight.

NUMA Awareness

On multi-socket systems, the scheduler discovers hardware topology via CpuTopology and uses AdaptiveBackoff (spin → yield → sleep) for NUMA-aware victim selection. A thief prefers victims on the same NUMA node to reduce cross-socket memory latency, and the backoff strategy ramps from a brief spin (good when work is about to appear) to a park (good when the system is fully idle).

Diffusion Pipeline Integration

Moirai is the execution substrate for every compute-heavy pass in the diffusion MRI pipeline. Each model’s fitting kernel dispatches across voxels through Moirai’s chunk operators:

The pattern is the same in every case:

1. A flat Vec<ModelResult> is allocated with one element per voxel (all elements MaybeUninit).
2. for_each_chunk_mut_with_state parcels the result array into chunks and assigns each to a worker.
3. Each worker initialises a scratch workspace once (init closure), then processes its assigned voxels sequentially within the chunk, reusing the workspace.
4. On completion, every slot is initialised and the Vec<MaybeUninit<T>> is safely reinterpreted as Vec<T>.

This is zero-copy (no intermediate per-voxel allocations), cache-friendly (consecutive voxels share the same signal array and scheme), and NUMA-aware (the scheduler prefers to assign chunks to workers on the same socket as the memory).

// Pattern for batched voxel-wise fitting (simplified)
fn fit_volume(
    signal: &[f64],
    scheme: &GradientScheme,
    config: &DtiConfig,
) -> Vec<DtiResult> {
    let n_voxels = signal.len() / scheme.len();
    let mut results: Vec<MaybeUninit<DtiResult>> = Vec::with_capacity(n_voxels);
    unsafe { results.set_len(n_voxels); }

    for_each_chunk_mut_with_state::<Adaptive, _, _, _, _>(
        &mut results,
        256,                           // voxels per task
        || DtiWorkspace::new(scheme),  // one design matrix per worker
        |ws, chunk| {
            for (i, slot) in chunk.iter_mut().enumerate() {
                let voxel_start = (ws.chunk_base + i) * scheme.len();
                slot.write(ws.fit(&signal[voxel_start..]));
            }
        },
    );

    // SAFETY: every slot written by the parallel region above.
    unsafe { std::mem::transmute(results) }
}

Performance Characteristics