5. Hybrid MPI + Rayon pmap
Model: MPI ranks, each with its own Rayon pool. Features: mpi,rayon,tenferro (or rsmpi-rt,rayon,tenferro). Example: mpi_mandelbrot_hybrid.
With both feature groups enabled, hataori::pmap is the hybrid entry point. The collective protocol is the same as in tutorial 3: the root owns the items and dynamically hands out batches. What changes is that each rank evaluates its batch inside a pooled Domain with the requested LocalMode, so in a one-rank-per-node deployment all cores of every node participate.
Initialize MPI for funneled threads
All MPI calls stay on the thread that calls pmap, while Rayon workers evaluate callbacks concurrently. That is exactly MPI_THREAD_FUNNELED, so ask for it explicitly and check what the implementation provided.
/// Hybrid `pmap` keeps every MPI call on the thread that initialized MPI
/// while Rayon workers evaluate callbacks, so MPI must be initialized with
/// at least `MPI_THREAD_FUNNELED`.
fn initialize_funneled() -> Universe {
let (universe, provided) =
mandelbrot_common::mpi_api::initialize_with_threading(Threading::Funneled)
.expect("MPI must not already be initialized or finalized");
assert!(
provided >= Threading::Funneled,
"hybrid pmap requires MPI_THREAD >= Funneled"
);
universe
}The collective call with a pooled domain
/// Every rank builds a Rayon domain and passes it to the collective. The
/// root hands out batches to itself and to remote ranks; each rank runs its
/// batch through its own pool with the requested `LocalMode`.
fn render<C: Communicator>(
world: &C,
param: &Param,
workers: usize,
batch_size: NonZeroUsize,
prefetch: bool,
) -> Option<Tensor> {
let rank = world.rank();
let (x, y) = param.make_axes();
// The worker count may differ per rank; the `PmapOptions` must not.
let cpu_set: Vec<usize> = (0..workers).collect();
let domain = Domain::managed(cpu_set, workers).expect("rayon domain");
let input = (rank == 0).then(|| (0..param.width).collect::<Vec<usize>>());
let columns = pmap(
world,
&domain,
PmapOptions {
batch_size,
// `Outer` spreads the columns of one batch over the pool.
local_mode: LocalMode::Outer,
// When true, each remote rank may hold one extra batch so that
// the next transfer overlaps the current computation.
prefetch,
..PmapOptions::default()
},
input,
|col_idx| Ok::<_, String>(mandelbrot_common::compute_column(param, x[col_idx], &y)),
)
.expect("pmap must succeed");
columns.map(|columns| {
mandelbrot_common::tensor_from_ordered_columns(param.width, param.height, columns)
})
}- Every rank builds a
Domain::managed(orDomain::external) pool; the worker count can differ per rank, but thePmapOptionsmust be identical. LocalMode::Outerspreads the columns of one batch over the pool, sobatch_sizeshould be at least the worker count.Inneris the choice when each callback is itself a parallel kernel, as in tutorial 2.- The callback and the values need the Rayon
Send/Syncbounds in addition to serde, but still not'static: the closure borrowsparam,x, andy. - The call must originate outside any Rayon worker; calling
pmapfrom inside a pool is rejected in preflight.
Batch size with a pool on every rank
/// Columns per `pmap` batch.
///
/// The batch size is `width / (size * workers * factor)`, clamped to at
/// least [`MIN_BATCH_COLUMNS`] columns and to at least `workers`, so that
/// `LocalMode::Outer` can spread one batch over the whole pool. The factor is
/// relative to the total number of compute threads (`size * workers`)
/// because in hybrid mode every batch is internally subdivided by Rayon.
fn batch_size(param: &Param, size: i32, workers: usize, factor: usize) -> NonZeroUsize {
let threads = (size as usize).max(1) * workers.max(1) * factor.max(1);
NonZeroUsize::new(
(param.width / threads)
.max(MIN_BATCH_COLUMNS)
.max(workers)
.min(param.width),
)
.expect("width is positive")
}The --batch-factor is now relative to the total number of compute threads (size × workers), because each batch is subdivided again by Rayon inside the rank.
Prefetch
The example renders the image twice, with prefetch: false and prefetch: true, and asserts that both give the identical image. With prefetch on, each remote rank may hold one additional batch, so receiving the next batch and sending the previous results overlap the current computation. The root’s own domain always stays at capacity one. The design and its bounds are described in P1 bounded prefetch.
Build and run
cargo build --release --no-default-features --features mpi,rayon,tenferro \
--example mpi_mandelbrot_hybrid
mpiexec -n 2 target/release/examples/mpi_mandelbrot_hybrid --workers 2 --width 1024 --height 1024On a laptop keep ranks × workers at or below the core count; --workers defaults to every logical core, which is the right value for one rank per node.
On Linux, Domain::managed pins each worker to a CPU from cpu_set and rejects CPUs outside the process’s affinity mask (DomainBuildError::CpuNotAllowed). Open MPI binds each rank to a single core by default for small jobs, which makes a two-worker managed domain fail. Either launch with mpiexec --bind-to none (or --map-by node:PE=<workers> to give each rank a core set), or use Domain::external with a pool the application built itself. scripts/check-tutorial-examples.sh passes --oversubscribe --bind-to none to Open MPI for this reason.
Next: 6. Runtime-loaded MPI.