nuvec-core 0.6.1

//! Graph evaluation.

use std::sync::Arc;

use crate::array::{Completion, Data, Node, NodeMap, NodeSet, State};
use crate::gpu::GpuEval;
use crate::kernels::{self, Operand};
use crate::primitive::Primitive;
use crate::tree::{Leaves, collect_leaves};
use crate::{Array, DeviceKind, Error, Result, view};

/// Evaluates the leaves of `tree` and everything they depend on, and waits for them. Each
/// node runs at most once, even across threads.
///
/// `tree` is any [`Leaves`] value: a slice or array of arrays (`eval(&[a, b])`), a `Vec`, a
/// model, or a borrowed view such as `&(&loss, &grads)`, which evaluates both without
/// cloning either.
///
/// Nodes on the GPU stream run as Metal kernels where one exists, and on the CPU (over the
/// same unified memory) otherwise. Returns the first error among `arrays`. Errors are
/// sticky: a failed array, and anything computed from it, fails with the same error on
/// every later evaluation or readback.
///
/// A failed GPU command buffer fails the evaluation that encoded it (its outputs and those
/// of its later command buffers) and every evaluation that read its results while they
/// were in flight; independent evaluations, earlier or later, are unaffected.
///
/// Several threads may evaluate arrays that share work: each node runs once, and a thread
/// that needs a node another thread is running waits for it (after submitting its own GPU
/// work), receiving its error if it fails.
pub fn eval<T: Leaves + ?Sized>(tree: &T) -> Result<()> {
    schedule(&collect_leaves(tree)).eval()
}

/// Starts evaluating the leaves of `tree` (any [`Leaves`] value, as [`eval`]) and returns
/// without waiting for GPU work: reading the arrays
/// (or evaluating them) later waits for exactly the work they need. As in MLX, use it to
/// overlap building the next computation with the GPU running this one.
///
/// CPU work runs on the calling thread before this returns. Errors are sticky: an error
/// already known is returned; a GPU error is reported by the first read or evaluation of
/// an affected array. Submission blocks while too many command buffers are in flight.
///
/// ```
/// use nuvec_core::{Result, array, async_eval};
///
/// # fn main() -> Result<()> {
/// let x = array![1.0, 2.0];
/// let y = &x * 3.0;
/// async_eval(&y)?;
/// // ... build the next computation here ...
/// assert_eq!(y.to_vec::<f32>()?, [3.0, 6.0]); // waits for `y` only
/// # Ok(())
/// # }
/// ```
pub fn async_eval<T: Leaves + ?Sized>(tree: &T) -> Result<()> {
    schedule(&collect_leaves(tree)).async_eval()
}

/// `arrays` with the work they need put in order, to be started later by
/// [`Scheduled::async_eval`] or [`Scheduled::eval`] with less delay: ordering a large
/// graph takes a while, which can then overlap other work, such as the GPU running the
/// previous step while the next one waits for its result.
///
/// The arrays are evaluated as [`async_eval`] or [`eval`] would at the time they are
/// started; nodes evaluated in between (by another evaluation) are skipped.
///
/// ```
/// use nuvec_core::{Result, array, schedule};
///
/// # fn main() -> Result<()> {
/// let x = array![1.0, 2.0];
/// let y = &(&x * 3.0) + 1.0;
/// let scheduled = schedule(std::slice::from_ref(&y));
/// // ... wait for something else here ...
/// scheduled.async_eval()?;
/// assert_eq!(y.to_vec::<f32>()?, [4.0, 7.0]);
/// # Ok(())
/// # }
/// ```
pub fn schedule(arrays: &[Array]) -> Scheduled {
    Scheduled {
        arrays: arrays.to_vec(),
        order: topological_order(arrays),
    }
}

/// Arrays whose evaluation is put in order but not started ([`schedule`]).
pub struct Scheduled {
    arrays: Vec<Array>,
    order: Vec<Arc<Node>>,
}

impl std::fmt::Debug for Scheduled {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Scheduled")
            .field("arrays", &self.arrays)
            .field("nodes", &self.order.len())
            .finish()
    }
}

impl Scheduled {
    /// The scheduled arrays.
    pub fn arrays(&self) -> &[Array] {
        &self.arrays
    }

    /// Starts evaluating the arrays without waiting for GPU work, as [`async_eval`].
    pub fn async_eval(self) -> Result<()> {
        let arrays = self.arrays.clone();
        self.submit().commit();
        match arrays.iter().find_map(Array::slot_error) {
            Some(err) => Err(err),
            None => Ok(()),
        }
    }

    /// Evaluates the arrays and waits for them, as [`eval`].
    pub fn eval(self) -> Result<()> {
        let arrays = self.arrays.clone();
        // Errors are reported through the arrays' slots below.
        let _ = self.submit().flush();
        wait(&arrays)
    }

    /// Encodes or computes every node in order; the caller commits (or flushes) the GPU
    /// work.
    fn submit(self) -> GpuEval {
        let mut gpu = GpuEval::new();
        for node in self.order {
            evaluate(&node, &mut gpu);
            // `node` is released here, so intermediates are freed as soon as their last
            // consumer has run (or, on the GPU, once their command buffer completes).
        }
        gpu
    }
}

/// Waits for `arrays`; the first error among them.
fn wait(arrays: &[Array]) -> Result<()> {
    for array in arrays {
        if let Err(err) = array.out().slot.wait() {
            return Err(err.clone());
        }
    }
    Ok(())
}

/// The widest breadth-first frontier when ordering a graph (MLX's `MLX_BFS_MAX_WIDTH`).
const BFS_MAX_WIDTH: usize = 20;

/// Unevaluated nodes reachable from `arrays`, inputs before consumers. Iterative, so graph
/// depth is limited only by memory.
///
/// As MLX's `eval_impl`: a tape is built breadth-first back from the outputs, a node
/// joining once all its consumers have, and run in reverse. Independent work (the query,
/// key and value projections of a layer, say) is then encoded side by side, so the GPU
/// runs it concurrently without barriers between. Past [`BFS_MAX_WIDTH`] pending nodes the
/// walk goes depth-first, which bounds how many intermediates are alive at once.
fn topological_order(arrays: &[Array]) -> Vec<Arc<Node>> {
    // Unevaluated nodes with their unevaluated inputs, depth-first.
    let mut inputs: NodeMap<*const Node, Vec<Arc<Node>>> = NodeMap::default();
    let mut done = NodeSet::default();
    let mut stack: Vec<Arc<Node>> = arrays.iter().map(|a| a.node.clone()).collect();
    while let Some(node) = stack.pop() {
        let key = Arc::as_ptr(&node);
        if inputs.contains_key(&key) || done.contains(&key) {
            continue;
        }
        let pending: Vec<Arc<Node>> = match &*node.lock() {
            State::Pending { inputs, .. } => inputs.iter().map(|a| a.node.clone()).collect(),
            // Running on another thread: its consumers wait for it (`input_data`).
            State::Running { .. } | State::Retained { .. } | State::Done | State::Failed { .. } => {
                done.insert(key);
                continue;
            }
        };
        stack.extend(pending.iter().cloned());
        inputs.insert(key, pending);
    }
    for list in inputs.values_mut() {
        list.retain(|n| !done.contains(&Arc::as_ptr(n)));
    }
    let roots: Vec<Arc<Node>> = arrays
        .iter()
        .map(|a| a.node.clone())
        .filter(|n| inputs.contains_key(&Arc::as_ptr(n)))
        .collect();
    // Consumers of each node (an input used twice counts twice; the outputs count once
    // each for the caller).
    let mut consumers: NodeMap<*const Node, usize> = NodeMap::default();
    for n in inputs.values().flatten().chain(&roots) {
        *consumers.entry(Arc::as_ptr(n)).or_default() += 1;
    }
    // The tape: `None` stands for the caller, whose inputs are `roots`.
    let mut tape: Vec<Option<Arc<Node>>> = vec![None];
    let mut deferred: Vec<(Option<Arc<Node>>, usize)> = Vec::new();
    let mut next = 0;
    while !consumers.is_empty() && (next < tape.len() || !deferred.is_empty()) {
        let (node, mut j) = if next < tape.len() {
            next += 1;
            (tape[next - 1].clone(), 0)
        } else {
            deferred.pop().expect("checked above")
        };
        let list = match &node {
            Some(n) => &inputs[&Arc::as_ptr(n)],
            None => &roots,
        };
        while j < list.len() {
            let key = Arc::as_ptr(&list[j]);
            if consumers.contains_key(&key) {
                if tape.len() - next >= BFS_MAX_WIDTH {
                    deferred.push((node.clone(), j));
                    break;
                }
                let left = consumers.get_mut(&key).expect("checked above");
                *left -= 1;
                if *left == 0 {
                    consumers.remove(&key);
                    tape.push(Some(list[j].clone()));
                }
            }
            j += 1;
        }
    }
    tape.into_iter().rev().flatten().collect()
}

/// Runs one node whose inputs are all evaluated, stores (or, for GPU work, defers) its
/// outputs, and detaches it (unless a running transform needs its graph).
///
/// The node's lock is held only to move it from `Pending` to `Running` and from `Running`
/// to its final state: computing (which may wait for inputs another thread is producing)
/// happens without it, so two evaluations sharing nodes never wait on each other's locks.
/// A node another thread is running is skipped here; its consumers wait for it in
/// [`input_data`].
fn evaluate(node: &Arc<Node>, gpu: &mut GpuEval) {
    let done = Arc::new(Completion::default());
    let (primitive, mut inputs) = {
        let mut state = node.lock();
        if !matches!(*state, State::Pending { .. }) {
            return; // evaluated, or running, on another thread
        }
        let running = State::Running {
            owner: std::thread::current().id(),
            done: done.clone(),
        };
        match std::mem::replace(&mut *state, running) {
            State::Pending { primitive, inputs } => (primitive, inputs),
            _ => unreachable!("checked above"),
        }
    };
    let running = Running {
        node,
        done,
        finished: false,
    };
    #[cfg(test)]
    counter::inc();
    let result = match donation(node, &primitive, &mut inputs, &mut *gpu) {
        Some(base) => in_place(node, &primitive, &inputs, base, gpu),
        None => compute(node, &primitive, &inputs, gpu),
    };
    let failure = match result {
        Ok((outputs, true)) => {
            gpu.defer(node.clone(), outputs);
            None
        }
        Ok((outputs, false)) => {
            for (output, data) in node.outputs.iter().zip(outputs) {
                let _ = output.slot.set(Ok(data));
            }
            None
        }
        Err(err) => {
            for output in &node.outputs {
                let _ = output.slot.set(Err(err.clone()));
            }
            Some(err)
        }
    };
    // A tracer keeps its graph while its transform runs, which releases it on exit.
    if crate::transforms::retain(node) {
        running.finish(State::Retained { primitive, inputs });
        return;
    }
    running.finish(match failure {
        Some(error) => State::Failed { error },
        None => State::Done,
    });
    drop(inputs); // detach: release the inputs outside the lock
}

/// A node this thread is running. Finishing it publishes its final state and wakes the
/// threads waiting for it; dropping it unfinished (a panic while computing) fails it, so
/// waiters never hang.
struct Running<'a> {
    node: &'a Arc<Node>,
    done: Arc<Completion>,
    finished: bool,
}

impl Running<'_> {
    fn finish(mut self, state: State) {
        *self.node.lock() = state;
        self.done.finish();
        // Dropping `self` now releases `done` without failing the node.
        self.finished = true;
    }
}

impl Drop for Running<'_> {
    fn drop(&mut self) {
        if self.finished {
            return;
        }
        let error = Error::Backend("the evaluation of this array panicked".into());
        for output in &self.node.outputs {
            let _ = output.slot.set(Err(error.clone()));
        }
        *self.node.lock() = State::Failed { error };
        self.done.finish();
    }
}

/// Waits until `node`, if another thread is running it, has finished. The owner may wait
/// for GPU work of ours in turn, so our open command buffer is committed first. Returns
/// the owner's error if the node failed.
fn wait_running(node: &Node, gpu: &mut GpuEval) -> Result<()> {
    let done = match &*node.lock() {
        State::Running { owner, done } => {
            if *owner == std::thread::current().id() {
                return Err(Error::InvalidArgument(
                    "eval: an array was evaluated from within its own evaluation".into(),
                ));
            }
            done.clone()
        }
        State::Failed { error } => return Err(error.clone()),
        _ => return Ok(()),
    };
    gpu.commit();
    done.wait();
    match &*node.lock() {
        State::Failed { error } => Err(error.clone()),
        _ => Ok(()),
    }
}

/// An input to be written in place (MLX's donation), and its buffer: an update's base
/// (`slice_update`, `norm_rope` or `paged_write` into a cache), or the input of an
/// element-wise operation or softmax whose buffer can take the output (same dtype, shape
/// and contiguous layout), as MLX's element-wise kernels reuse their inputs.
fn donation(
    node: &Node,
    primitive: &Primitive,
    inputs: &mut [Array],
    gpu: &mut GpuEval,
) -> Option<(usize, Data)> {
    let device = node.stream.device;
    if node.is_tracer() {
        return None;
    }
    if device.kind == DeviceKind::Cpu {
        // MLX's CPU donation for updates: the base's buffer, if nothing else holds it.
        return match primitive {
            Primitive::SliceUpdate { .. } if donatable_size(inputs[0].size()) => {
                owned_on_cpu(&mut inputs[0]).map(|data| (0, data))
            }
            _ => None,
        };
    }
    if !gpu.available(device) {
        return None;
    }
    let base = match primitive {
        Primitive::SliceUpdate { .. }
        | Primitive::Fast(crate::fast::imp::FastOp::NormRope { cache: true, .. })
        | Primitive::Fast(crate::fast::imp::FastOp::PagedWrite) => Some(0),
        Primitive::Fast(crate::fast::imp::FastOp::GatedDeltaSlots) => Some(5),
        Primitive::Fast(crate::fast::imp::FastOp::ConvSlots { .. }) => Some(2),
        // MLX's donation for an all-reduce or send: the worker writes the input's buffer.
        Primitive::Collective {
            op:
                crate::distributed::CollectiveOp::Reduce(_)
                | crate::distributed::CollectiveOp::Send { .. },
            ..
        } => Some(0),
        _ => None,
    };
    if let Some(at) = base {
        if !donatable_size(inputs[at].size()) {
            return None;
        }
        return owned(&mut inputs[at], gpu).map(|data| (at, data));
    }
    let [out] = node.outputs.as_slice() else {
        return None;
    };
    for i in in_place_inputs(primitive, inputs, out) {
        if let Some(data) = owned(&mut inputs[i], gpu) {
            return Some((i, data));
        }
    }
    None
}

/// Whether an update base of `size` elements may be written in place: the update kernels
/// index with 32 bits (as the GPU's `supports` requires of their operands).
fn donatable_size(size: usize) -> bool {
    size > 0 && size < u32::MAX as usize
}

/// The inputs of an element-wise operation or softmax whose buffer the output may take:
/// the output's dtype (float, where every such kernel exists on the GPU) and shape.
fn in_place_inputs(
    primitive: &Primitive,
    inputs: &[Array],
    out: &crate::array::Output,
) -> Vec<usize> {
    use crate::{DType, UnaryOp};
    let candidates: &[usize] = match primitive {
        Primitive::Unary(op) if *op != UnaryOp::Conjugate => &[0],
        Primitive::Binary(_) => &[0, 1],
        Primitive::Select => &[1, 2],
        Primitive::Fast(
            crate::fast::imp::FastOp::Softmax | crate::fast::imp::FastOp::CausalSoftmax { .. },
        ) => &[0],
        _ => &[],
    };
    let float = matches!(out.dtype, DType::F16 | DType::BF16 | DType::F32);
    if !float || out.size == 0 {
        return Vec::new();
    }
    // From 2^32 elements, only layouts the GPU's large kernels address (as `supports`).
    if out.size >= u32::MAX as usize {
        let large = match primitive {
            Primitive::Fast(_) => {
                let axis = out.shape.last().copied().unwrap_or(1).max(1);
                out.size / axis * 1024 < u32::MAX as usize
            }
            _ => {
                let metas: Vec<&crate::array::Output> = inputs.iter().map(Array::out).collect();
                crate::gpu::large_elementwise(&metas, &out.shape)
            }
        };
        if !large {
            return Vec::new();
        }
    }
    candidates
        .iter()
        .copied()
        .filter(|&i| inputs[i].dtype() == out.dtype && inputs[i].dims() == out.shape.as_slice())
        .collect()
}

/// The buffer of `input`, taken over to be written in place: only when the evaluation owns
/// it outright (its node has no other strong handle, so no other array, view or consumer
/// refers to it) and its data is the whole of a writable buffer with no other handle.
/// Ownership proves the write exclusive, so no reference counting of other users is
/// needed.
///
/// The input may be finished, still in flight in an earlier evaluation (pending work holds
/// nodes and data only weakly, so the in-flight copy is then the data's only handle;
/// taking it over under the lock that publishing takes means it is never published), or
/// computed earlier in this evaluation. The buffer's fences and the encoder's barriers
/// order the write after the work that produces and reads it.
///
/// Invariant (concurrency): the `Arc::strong_count` checks are race-free only because a
/// node no other handle refers to is out of every graph: another thread can reach a node
/// only through an `Arc` (an array, a consumer's inputs, or its own evaluation order),
/// and each of those counts. Weak references from committed GPU work upgrade only to
/// publish, under the in-flight lock that donation takes too. Arrays put in order by
/// [`schedule`] hold their nodes strongly until started, so scheduled-but-unstarted work
/// disables donation of what it reads.
fn owned(input: &mut Array, gpu: &mut GpuEval) -> Option<Data> {
    // Output `at` of the node (another output, if any, has no holder either: the node's
    // only strong handle is this input).
    let at = input.output as usize;
    let base = &mut input.node;
    let usable = |data: &Data, out: &crate::array::Output, handles: usize| {
        out.flags.row_contiguous
            && crate::gpu::writable(data, out)
            && Arc::strong_count(&data.buffer) == handles
    };
    // Produced earlier in this evaluation: taken over from it if nothing else refers to it.
    if let Some(out) = base.outputs.get(at)
        && let Some(data) = gpu.take_pending(base, at, |data, handles| usable(data, out, handles))
    {
        return Some(data);
    }
    // The in-flight cell is locked first: publishing locks it before it upgrades the node
    // and drops the upgraded handle before it unlocks, so the count below is the node's real
    // handles.
    let cell = base.outputs.get(at)?.inflight.clone();
    let mut state = cell.lock();
    if Arc::strong_count(base) != 1 {
        return None;
    }
    // Finished: no pending work refers to it any more, so it is owned outright.
    if let Some(base) = Arc::get_mut(base) {
        let out = base.outputs.get_mut(at)?;
        if let Some(Ok(data)) = out.slot.get() {
            // The slot and (if it was published from in-flight work) the in-flight copy.
            let published = matches!(*state, crate::array::InFlightState::Set(_));
            if !usable(data, out, 1 + usize::from(published)) {
                return None;
            }
            if published {
                *state = crate::array::InFlightState::Unset;
            }
            return match out.slot.take() {
                Some(Ok(data)) => Some(data),
                _ => None,
            };
        }
    }
    // In flight: only weak references from pending work remain besides ours.
    let out = base.outputs.get(at)?;
    if out.slot.get().is_some() {
        return None;
    }
    let crate::array::InFlightState::Set(data) = &*state else {
        return None;
    };
    if !usable(data, out, 1) {
        return None;
    }
    match std::mem::replace(&mut *state, crate::array::InFlightState::Donated) {
        crate::array::InFlightState::Set(data) => Some(data),
        _ => None,
    }
}

/// A finished input the CPU may write in place: no other handle to its node or buffer, the
/// buffer holds exactly the row-major array, and it is CPU-writable (not a mapped file, no
/// GPU work using it).
fn owned_on_cpu(input: &mut Array) -> Option<Data> {
    let at = input.output as usize;
    let base = Arc::get_mut(&mut input.node)?;
    let out = base.outputs.get_mut(at)?;
    if out.inflight.get().is_some() {
        return None;
    }
    let Some(Ok(data)) = out.slot.get() else {
        return None;
    };
    let whole = data.offset == 0 && data.buffer.len() == out.size * out.dtype.size();
    if !(out.flags.row_contiguous
        && whole
        && Arc::strong_count(&data.buffer) == 1
        && data.buffer.cpu_writable())
    {
        return None;
    }
    match out.slot.take() {
        Some(Ok(data)) => Some(data),
        _ => None,
    }
}

/// An update written into its donated base (input 0), or an element-wise operation or
/// softmax written into its donated input `at`.
fn in_place(
    node: &Node,
    primitive: &Primitive,
    inputs: &[Array],
    (at, base): (usize, Data),
    gpu: &mut GpuEval,
) -> Result<(Vec<Data>, bool)> {
    let (mut pending, mut foreign) = (false, Vec::new());
    let mut rest = Vec::with_capacity(inputs.len());
    for (i, a) in inputs.iter().enumerate() {
        rest.push(match i == at {
            true => None,
            false => Some(input_data(a, gpu, &mut pending, &mut foreign)?),
        });
    }
    let out = &node.outputs[0];
    if let Primitive::Collective { op, group } = primitive {
        let outputs = gpu.collective(op, *group, inputs, Vec::new(), out, Some(base))?;
        return Ok((outputs, true));
    }
    let operands = |skip: usize| -> Vec<Operand<'_>> {
        inputs
            .iter()
            .zip(&rest)
            .skip(skip)
            .filter_map(|(a, d)| d.as_ref().map(|d| (a.out(), d)))
            .collect()
    };
    if let Primitive::Fast(crate::fast::imp::FastOp::GatedDeltaSlots) = primitive {
        let pool = inputs[at].out();
        let all = operands(0);
        let datas = gpu.gated_delta_slots_in_place(&all, pool, base, &node.outputs)?;
        return Ok((datas, true));
    }
    if let Primitive::Fast(crate::fast::imp::FastOp::ConvSlots { silu }) = primitive {
        let pool = inputs[at].out();
        let all = operands(0);
        let datas = gpu.conv_slots_in_place(&all, pool, base, &node.outputs, *silu)?;
        return Ok((datas, true));
    }
    let data = match primitive {
        Primitive::Fast(crate::fast::imp::FastOp::PagedWrite) => {
            gpu.paged_write_in_place(base, &operands(1), out)?
        }
        Primitive::Fast(crate::fast::imp::FastOp::Softmax) => {
            gpu.softmax_in_place(base, out, None)?
        }
        Primitive::Fast(crate::fast::imp::FastOp::CausalSoftmax { rows }) => {
            gpu.softmax_in_place(base, out, Some(*rows as usize))?
        }
        Primitive::Fast(op) => gpu.norm_rope_in_place(base, op, &operands(1), out)?,
        Primitive::SliceUpdate { .. } if node.stream.device.kind == DeviceKind::Cpu => {
            // The CPU reads the update in place: wait for the GPU work that produces it.
            for input in &foreign {
                input.out().slot.wait().clone()?;
            }
            if pending {
                gpu.flush()?;
            }
            let slices = primitive.dim_slices(out.shape.len())?;
            let mut base = base;
            let buffer = Arc::get_mut(&mut base.buffer).expect("donated: the only handle");
            crate::kernels::slice_update_into(buffer, &out.shape, operands(1)[0], &slices);
            return Ok((vec![base], false));
        }
        Primitive::SliceUpdate { .. } => {
            let slices = primitive.dim_slices(out.shape.len())?;
            gpu.slice_update_in_place(base, out, operands(1)[0], &slices)?
        }
        _ => {
            let metas: Vec<&crate::array::Output> = inputs.iter().map(Array::out).collect();
            let datas: Vec<Option<&Data>> = rest.iter().map(Option::as_ref).collect();
            gpu.elementwise_in_place(primitive, &metas, &datas, base, out)?
        }
    };
    Ok((vec![data], true))
}

/// An input's data: from this evaluation's GPU work, evaluated, or in flight in another
/// evaluation (then `pending`, and recorded in `foreign`).
fn input_data<'a>(
    input: &'a Array,
    gpu: &mut GpuEval,
    pending: &mut bool,
    foreign: &mut Vec<&'a Array>,
) -> Result<Data> {
    if let Some(d) = gpu.lookup(&input.node, input.output) {
        *pending = true;
        Ok(d)
    } else if let Some(result) = input.out().slot.get() {
        result.clone()
    } else if let Some(d) = input.out().inflight.get() {
        // Committed by another evaluation on the device timeline: GPU work encoded now
        // runs after it, so it may read the data at once.
        *pending = true;
        foreign.push(input);
        gpu.read_in_flight(input);
        Ok(d)
    } else {
        // Produced by another evaluation (another thread, or an earlier `async_eval`):
        // submit our own GPU work first, so two evaluations never wait on each other, then
        // wait for the other thread to finish running the node.
        gpu.commit();
        wait_running(&input.node, gpu)?;
        if let Some(result) = input.out().slot.get() {
            return result.clone();
        }
        if let Some(d) = input.out().inflight.get() {
            *pending = true;
            foreign.push(input);
            gpu.read_in_flight(input);
            return Ok(d);
        }
        // Deferred in the producer's open GPU work, which it commits before it waits for
        // anything and at the end of its evaluation.
        input.out().slot.wait().clone()
    }
}

/// Computes a node's outputs. The flag is true if they depend on GPU work that has not
/// completed yet (and must be deferred until it has).
fn compute(
    node: &Node,
    primitive: &Primitive,
    inputs: &[Array],
    gpu: &mut GpuEval,
) -> Result<(Vec<Data>, bool)> {
    let mut data = Vec::with_capacity(inputs.len());
    let mut pending = false;
    // Inputs read from another evaluation's in-flight GPU work.
    let mut foreign = Vec::new();
    for input in inputs {
        data.push(input_data(input, gpu, &mut pending, &mut foreign)?);
    }
    let operands: Vec<Operand<'_>> = inputs
        .iter()
        .zip(&data)
        .map(|(a, d)| (a.out(), d))
        .collect();
    // Waits until foreign in-flight inputs are complete (for work that does not run on the
    // device timeline).
    let await_foreign = || -> Result<()> {
        for input in &foreign {
            input.out().slot.wait().clone()?;
        }
        Ok(())
    };
    if matches!(primitive, Primitive::Depends { .. }) {
        // An ordering point (blocked attention, checkpoints): the work before it is
        // encoded, so the intermediates it no longer needs return to the allocator now,
        // whatever their size, before the work after it allocates.
        gpu.release_unreferenced();
    }
    if let Some(views) = view_outputs(primitive, &operands)? {
        // Views do no work: they are published with this evaluation's GPU work (which runs
        // after any foreign producer) or, without GPU work, once the producers complete.
        gpu.depend_on(&foreign);
        return Ok((views, pending));
    }
    let device = node.stream.device;
    if device.kind != DeviceKind::Cpu {
        if !gpu.available(device) {
            return Err(Error::UnsupportedDevice(device));
        }
        if gpu.supports(primitive, &operands, &node.outputs) {
            return Ok((gpu.run(primitive, &operands, &node.outputs)?, true));
        }
        if let Primitive::Collective { op, group } = primitive {
            // Other evaluations' in-flight inputs: the GPU copy is ordered after them.
            gpu.depend_on(&foreign);
            let outputs = gpu.collective(op, *group, inputs, data, &node.outputs[0], None)?;
            return Ok((outputs, true));
        }
    }
    // Collectives in the open command buffer wait for its signals: submit it before
    // anything here may block (a collective on the CPU waits for the queued ones).
    if gpu.collectives_pending() {
        gpu.commit();
    }
    // The CPU reads inputs in place: wait for the GPU work that produces them.
    await_foreign()?;
    if pending {
        gpu.flush()?;
    }
    Ok((cpu(primitive, &operands)?, false))
}

/// Outputs of view primitives, which share their input's buffer (on every device), or
/// `None` for primitives that compute.
fn view_outputs(primitive: &Primitive, operands: &[Operand<'_>]) -> Result<Option<Vec<Data>>> {
    let share = |offset| {
        Some(vec![Data {
            buffer: operands[0].1.buffer.clone(),
            offset,
        }])
    };
    Ok(match primitive {
        Primitive::Transpose(_)
        | Primitive::Broadcast(_)
        | Primitive::StopGradient
        | Primitive::Depends { .. } => share(operands[0].1.offset),
        Primitive::Slice { .. } => {
            let (meta, data) = operands[0];
            let (_, delta) =
                view::slice_geometry(&meta.strides, &primitive.dim_slices(meta.shape.len())?);
            share(data.offset + delta)
        }
        Primitive::Custom(call) => Some(
            operands[call.primals..]
                .iter()
                .map(|(_, data)| (*data).clone())
                .collect(),
        ),
        Primitive::Reshape(shape) => {
            let (meta, data) = operands[0];
            view::nocopy_reshape(&meta.shape, &meta.strides, shape).and_then(|_| share(data.offset))
        }
        // Removing or inserting size-1 dimensions is always a view.
        Primitive::Squeeze(_) | Primitive::ExpandDims(_) => share(operands[0].1.offset),
        // The offset counts elements: in the new dtype it must stay whole.
        Primitive::View(dtype) => {
            let (meta, data) = operands[0];
            let (from, to) = (meta.dtype.size() as isize, dtype.size() as isize);
            let bytes = data.offset * from;
            let shared = from == to || (meta.flags.row_contiguous && bytes % to == 0);
            shared.then(|| share(bytes / to)).flatten()
        }
        _ => None,
    })
}

/// Computing primitives on the CPU.
fn cpu(primitive: &Primitive, operands: &[Operand<'_>]) -> Result<Vec<Data>> {
    let new = |buffer| Ok(vec![Data::new(buffer)]);
    match primitive {
        Primitive::Unary(op) => new(kernels::unary(*op, operands[0])?),
        Primitive::Binary(op) => new(kernels::binary(*op, operands[0], operands[1])?),
        Primitive::Select => new(kernels::select(operands[0], operands[1], operands[2])?),
        Primitive::AsType(dtype) => new(kernels::cast(operands[0], *dtype)?),
        Primitive::ToFp8 => new(kernels::to_fp8(operands[0])?),
        Primitive::SortedOffsets { groups } => new(kernels::sorted_offsets(operands[0], *groups)?),
        Primitive::FromFp8(dtype) => new(kernels::from_fp8(operands[0], *dtype)?),
        Primitive::Reduce { op, axes } => new(kernels::reduce(*op, operands[0], axes)?),
        Primitive::ArgReduce { op, axis } => new(kernels::arg_reduce(*op, operands[0], *axis)?),
        Primitive::GatherAxis { axis } => {
            new(kernels::gather_axis(operands[0], operands[1], *axis)?)
        }
        Primitive::ScatterAxis { axis, reduce } => new(kernels::scatter_axis(
            operands[0],
            operands[1],
            operands[2],
            *axis,
            *reduce,
        )?),
        Primitive::Concatenate { axis } => {
            let mut shape = operands[0].0.shape.clone();
            shape[*axis] = operands.iter().map(|(m, _)| m.shape[*axis]).sum();
            new(kernels::concatenate(operands, &shape, *axis)?)
        }
        Primitive::Sort { axis } => new(kernels::sort(operands[0], *axis, false)?),
        Primitive::ArgSort { axis } => new(kernels::sort(operands[0], *axis, true)?),
        Primitive::Scan {
            op,
            axis,
            reverse,
            inclusive,
        } => new(kernels::scan(
            *op,
            operands[0],
            *axis,
            *reverse,
            *inclusive,
        )?),
        Primitive::Matmul => new(kernels::matmul(operands[0], operands[1])?),
        Primitive::AddMM { .. } => {
            let (alpha, beta) = primitive.addmm_factors();
            new(kernels::addmm(
                operands[0],
                operands[1],
                operands[2],
                alpha,
                beta,
            )?)
        }
        Primitive::RandomBits { shape, .. } => {
            let n: usize = shape.iter().product();
            new(kernels::random_bits(
                operands[0],
                n * primitive_width(primitive),
            )?)
        }
        Primitive::Fast(op) => op.eval_cpu(operands),
        Primitive::MetalKernel(_) => Err(Error::InvalidArgument(
            "metal_kernel: custom kernels run on the GPU only".into(),
        )),
        Primitive::Compiled(fused) => fused.eval_cpu(operands),
        Primitive::Collective { op, group } => crate::distributed::eval_cpu(op, *group, operands),
        Primitive::Quantize { quant } => {
            let (group_size, bits) = quant.checked("quantize")?;
            if quant.mode != crate::QuantMode::Affine {
                let (packed, scales) = kernels::fp_quantize(operands[0], group_size, bits)?;
                return Ok(vec![Data::new(packed), Data::new(scales)]);
            }
            let (packed, scales, biases) = kernels::quantize(operands[0], group_size, bits)?;
            Ok(vec![
                Data::new(packed),
                Data::new(scales),
                Data::new(biases),
            ])
        }
        Primitive::Dequantize { quant, dtype } => {
            let (group_size, bits) = quant.checked("dequantize")?;
            new(match quant.mode {
                crate::QuantMode::Affine => {
                    kernels::dequantize(operands[0], operands[1], operands[2], group_size, bits)?
                }
                _ => kernels::fp_dequantize(operands[0], operands[1], group_size, bits, *dtype)?,
            })
        }
        Primitive::QuantizedMatmul { quant, transpose } => {
            let (group_size, bits) = quant.checked("quantized_matmul")?;
            new(kernels::quantized_matmul(
                operands, *quant, group_size, bits, *transpose,
            )?)
        }
        Primitive::SliceUpdate { .. } => new(kernels::slice_update(
            operands[0],
            operands[1],
            &primitive.dim_slices(operands[0].0.shape.len())?,
        )?),
        // A reshape that `view_outputs` could not express as a view copies; so does a `view`,
        // whose copy is then read in the new dtype.
        Primitive::Reshape(_)
        | Primitive::Squeeze(_)
        | Primitive::ExpandDims(_)
        | Primitive::Copy
        | Primitive::View(_) => new(kernels::contiguous(operands[0].0, operands[0].1)?),
        Primitive::Transpose(_)
        | Primitive::Broadcast(_)
        | Primitive::Slice { .. }
        | Primitive::StopGradient
        | Primitive::Depends { .. }
        | Primitive::Custom(_) => {
            unreachable!("views are handled by view_outputs")
        }
        #[cfg(test)]
        Primitive::AddMultiply => {
            use crate::primitive::BinaryOp;
            Ok(vec![
                Data::new(kernels::binary(BinaryOp::Add, operands[0], operands[1])?),
                Data::new(kernels::binary(
                    BinaryOp::Multiply,
                    operands[0],
                    operands[1],
                )?),
            ])
        }
    }
}

fn primitive_width(primitive: &Primitive) -> usize {
    match primitive {
        Primitive::RandomBits { width, .. } => usize::from(*width),
        _ => unreachable!(),
    }
}

/// Per-thread count of evaluated nodes, for run-once tests.
#[cfg(test)]
pub(crate) mod counter {
    use std::cell::Cell;

    thread_local!(static EVALUATED: Cell<usize> = const { Cell::new(0) });

    pub(crate) fn inc() {
        EVALUATED.set(EVALUATED.get() + 1);
    }
    pub(crate) fn take() -> usize {
        EVALUATED.replace(0)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::buffer::live;

    #[test]
    fn shared_node_runs_once() {
        let x = Array::from_slice(&[1.0f32, 2.0], [2]);
        let shared = &x + &x;
        let left = &shared * &x;
        let right = &shared + &x;
        let top = &left + &right;
        counter::take();
        eval(&[top.clone(), left.clone()]).unwrap();
        assert_eq!(counter::take(), 4, "shared, left, right, top: once each");
        top.eval().unwrap();
        assert_eq!(counter::take(), 0, "already evaluated");
        assert_eq!(top.to_vec::<f32>().unwrap(), vec![5.0, 14.0]);
    }

    #[test]
    fn running_releases_its_completion() {
        // Both ends of a running node release the shared `Completion`: finishing it and
        // dropping it unfinished (a panic while computing).
        let x = Array::from_slice(&[1.0f32], [1]);
        for finish in [true, false] {
            let done = Arc::new(Completion::default());
            let running = Running {
                node: &x.node,
                done: done.clone(),
                finished: false,
            };
            if finish {
                running.finish(State::Done);
            } else {
                drop(running);
            }
            assert_eq!(Arc::strong_count(&done), 1, "finish = {finish}");
        }
    }

    #[test]
    fn intermediates_freed() {
        // The live count is per thread; GPU buffers are released on Metal's completion
        // thread, so this checks the CPU path.
        crate::set_default_stream(crate::Stream::default_for(crate::Device::cpu()));
        let before = live::count();
        let mut x = Array::from(0i64);
        let one = Array::from(1i64);
        let steps = if cfg!(miri) { 100 } else { 10_000 };
        for _ in 0..steps {
            x = &x + &one;
        }
        x.eval().unwrap();
        drop(one);
        assert_eq!(
            live::count() - before,
            1,
            "only the result's buffer is alive"
        );
        assert_eq!(x.item::<i64>().unwrap(), steps as i64);
    }

    #[test]
    fn multi_output_once() {
        let a = Array::from_slice(&[2i32, 3], [2]);
        let b = Array::from_slice(&[5i32, 7], [2]);
        let mut outs = Array::from_op(Primitive::AddMultiply, vec![a, b]).unwrap();
        let product = outs.pop().unwrap();
        let sum = outs.pop().unwrap();
        counter::take();
        eval(&[sum.clone(), product.clone()]).unwrap();
        assert_eq!(counter::take(), 1);
        drop(sum);
        assert_eq!(product.to_vec::<i32>().unwrap(), vec![10, 21]);
    }

    #[test]
    fn slice_update_writes_a_strided_region() {
        // Base: a transposed view; region: rows 0..3 step 2, columns reversed.
        let base = crate::ops::imp::arange(0.0, 12.0, 1.0, crate::DType::I32)
            .unwrap()
            .try_reshape(&[4, 3])
            .unwrap()
            .t(); // [[0,3,6,9],[1,4,7,10],[2,5,8,11]]
        let update = Array::from_slice(&[-1i32, -2, -3, -4, -5, -6, -7, -8], [2, 4]);
        let primitive = Primitive::SliceUpdate {
            start: vec![0, 3],
            step: vec![2, -1],
            len: vec![2, 4],
        };
        let out = Array::from_op(primitive.clone(), vec![base.clone(), update.clone()])
            .unwrap()
            .remove(0);
        assert_eq!(
            out.to_vec::<i32>().unwrap(),
            [-4, -3, -2, -1, 1, 4, 7, 10, -8, -7, -6, -5]
        );
        let wrong = Array::from_slice(&[0i32; 3], [3]);
        assert!(Array::from_op(primitive, vec![base, wrong]).is_err());
    }

    #[test]
    fn tracer_nodes_keep_their_graph_only_while_tracing() {
        let x = Array::from(2.0f32);
        let kept = std::cell::Cell::new(false);
        let g = crate::grad(|x: &Array| {
            let y = x * x;
            y.eval()?;
            kept.set(matches!(&*y.node.lock(), State::Retained { .. }));
            Ok(y)
        });
        assert_eq!(g(&x).unwrap().item::<f32>().unwrap(), 4.0);
        assert!(kept.get(), "retained inside a transform");
        let traced = crate::vjp_flat(|xs| Ok(vec![&xs[0] * &xs[0]]), &[x], &[Array::from(1.0f32)])
            .unwrap()
            .0
            .remove(0);
        traced.eval().unwrap();
        assert!(
            matches!(&*traced.node.lock(), State::Done),
            "outside a transform"
        );
    }

    #[test]
    fn huge_buffers_are_never_donated() {
        // The rule both donation paths (element-wise inputs and update bases) apply.
        assert!(donatable_size(1));
        assert!(donatable_size(u32::MAX as usize - 1));
        assert!(!donatable_size(u32::MAX as usize));
        assert!(!donatable_size(1 << 33));
        assert!(!donatable_size(0));
        // A lazy update of a base of 2^32 elements (metadata only) is never written in
        // place. (A single dimension cannot exceed `i32::MAX`, so the base is 2-D.)
        let base =
            crate::ops::imp::broadcast_to(&Array::from(0.0f32), &[1 << 16, 1 << 16]).unwrap();
        let update = Array::from_slice(&[1.0f32], [1, 1]);
        let primitive = Primitive::SliceUpdate {
            start: vec![0, 0],
            step: vec![1, 1],
            len: vec![1, 1],
        };
        let out = Array::from_op(primitive.clone(), vec![base.clone(), update.clone()])
            .unwrap()
            .remove(0);
        let mut inputs = vec![base, update];
        let mut gpu = GpuEval::new();
        assert!(donation(&out.node, &primitive, &mut inputs, &mut gpu).is_none());
    }

    #[test]
    fn evaluation_detaches_inputs() {
        let a = Array::from(1u8);
        let c = &a + &a;
        assert!(c.primitive_t().is_some());
        c.eval().unwrap();
        assert_eq!(c.primitive_t(), None);
        assert_eq!(Arc::strong_count(&a.node), 1, "c no longer references a");
    }

    /// Positions of nodes in the evaluation order of `arrays`, checked to be topological.
    fn positions(arrays: &[Array]) -> NodeMap<*const Node, usize> {
        let order = topological_order(arrays);
        let at: NodeMap<*const Node, usize> = order
            .iter()
            .enumerate()
            .map(|(i, n)| (Arc::as_ptr(n), i))
            .collect();
        assert_eq!(at.len(), order.len(), "each node once");
        for (i, node) in order.iter().enumerate() {
            if let State::Pending { inputs, .. } = &*node.lock() {
                for input in inputs {
                    if let Some(&j) = at.get(&Arc::as_ptr(&input.node)) {
                        assert!(j < i, "inputs come first");
                    }
                }
            }
        }
        at
    }

    #[test]
    fn independent_work_is_ordered_side_by_side() {
        let x = Array::from_slice(&[1.0f32, 2.0], [2]);
        let (a1, b1) = (&x + 1.0, &x + 3.0);
        let (a2, b2) = (&a1 * 2.0, &b1 * 4.0);
        let top = &a2 + &b2;
        let at = positions(std::slice::from_ref(&top));
        let p = |a: &Array| at[&Arc::as_ptr(&a.node)];
        assert!(
            p(&a1).max(p(&b1)) < p(&a2).min(p(&b2)),
            "both chains advance together"
        );
        assert_eq!(top.to_vec::<f32>().unwrap(), [20.0, 26.0]);
    }

    #[test]
    fn wide_and_shared_graphs_are_ordered() {
        let x = Array::from_slice(&[1i32], [1]);
        let one = Array::from_slice(&[1i32], [1]);
        let shared = &x + &x;
        // 100 chains of different depths (wider than the breadth-first limit), joined.
        let mut joined = shared.clone();
        for i in 0..100 {
            let mut c = &shared + &Array::from_slice(&[i], [1]);
            for _ in 0..i % 7 {
                c = &c * &one;
            }
            joined = &joined + &c;
        }
        positions(&[joined.clone(), shared.clone(), joined.clone()]);
        counter::take();
        eval(&[joined.clone(), shared.clone()]).unwrap();
        let nodes = 1 + 100 * 2 + (0..100).map(|i| i % 7).sum::<usize>();
        assert_eq!(counter::take(), nodes, "every node once");
        assert_eq!(
            joined.item::<i32>().unwrap(),
            2 + 100 * 2 + (0..100).sum::<i32>()
        );
    }

    #[test]
    fn scheduled_evaluation_skips_work_done_meanwhile() {
        let x = Array::from_slice(&[1.0f32, 2.0], [2]);
        let shared = &x * &x;
        let top = &shared + &x;
        let scheduled = schedule(std::slice::from_ref(&top));
        assert_eq!(scheduled.arrays().len(), 1);
        counter::take();
        shared.eval().unwrap();
        scheduled.eval().unwrap();
        assert_eq!(counter::take(), 2, "shared once, then top");
        assert_eq!(top.to_vec::<f32>().unwrap(), [2.0, 6.0]);
    }
}