moirai-parallel 0.7.0

//! Whole-unit tasks over `K` same-type buffers, split in lockstep: the fused
//! alternative to one pass per output.

use super::driver::drive_unit_tasks;
use crate::ops::shards::BufferArray;
use crate::policy::ExecutionPolicy;

/// [`crate::for_each_unit_task_mut_with`] over `K` buffers of one type, split in
/// lockstep.
///
/// The pair and triple forms take their buffers by separate type, which suits
/// two or three outputs of different types. A pass whose outputs share a type
/// and whose count is a property of the problem -- the six fields a
/// velocity-Verlet kick-and-drift writes, say -- takes them as one array
/// instead, and stays one pass rather than becoming `K` of them. Splitting a
/// fused pass into one call per output multiplies the parallel regions, and
/// each region's wake costs more than the traffic such a fusion saves.
///
/// Every buffer must have the same length, and that length must be whole
/// units. `f` receives each task's runs in the order the buffers were given.
///
/// # Examples
///
/// ```
/// use moirai_parallel::{for_each_unit_task_many_mut_with, WorkBytes};
///
/// let (mut xs, mut ys) = (vec![0_u64; 8], vec![1_u64; 8]);
/// for_each_unit_task_many_mut_with::<WorkBytes<{ 1 << 20 }>, _, _, _, _, 2>(
///     [&mut xs, &mut ys],
///     4,
///     64,
///     || (),
///     |(), _first_unit, [xs, ys]| {
///         for (x, y) in xs.iter_mut().zip(ys.iter_mut()) {
///             *x += 2;
///             *y += *x;
///         }
///     },
/// );
/// assert!(xs.iter().all(|&x| x == 2) && ys.iter().all(|&y| y == 3));
/// ```
#[track_caller]
pub fn for_each_unit_task_many_mut_with<P, T, S, Init, F, const K: usize>(
    outputs: [&mut [T]; K],
    unit_len: usize,
    unit_bytes: usize,
    init: Init,
    f: F,
) where
    P: ExecutionPolicy,
    T: Send,
    Init: Fn() -> S + Send + Sync,
    F: Fn(&mut S, usize, [&mut [T]; K]) + Send + Sync,
{
    let Some(first) = outputs.first() else {
        return;
    };
    let len = first.len();
    for buffer in &outputs {
        assert_eq!(
            buffer.len(),
            len,
            "many unit tasks need buffers of one length",
        );
    }
    drive_unit_tasks::<P, _, _, _, _>(
        BufferArray { buffers: outputs },
        unit_len,
        unit_bytes,
        "moirai global executor: for_each_unit_task_many_mut_with",
        init,
        f,
    );
}