nuvec_core/
eval.rs

1//! Graph evaluation.
2
3use std::collections::HashSet;
4use std::sync::Arc;
5
6use crate::array::{Data, Node, State};
7use crate::kernels::{self, Operand};
8use crate::primitive::Primitive;
9use crate::{Array, DeviceKind, Error, Result, view};
10
11/// Evaluates `arrays` and everything they depend on. Each node runs at most once, even
12/// across threads.
13///
14/// Returns the first error among `arrays`. Errors are sticky: a failed array, and anything
15/// computed from it, fails with the same error on every later evaluation or readback.
16pub fn eval(arrays: &[Array]) -> Result<()> {
17    for node in topological_order(arrays) {
18        evaluate(&node);
19        // `node` is released here, so intermediates are freed as soon as their last
20        // consumer has run.
21    }
22    match arrays
23        .iter()
24        .find_map(|a| a.out().slot.get().and_then(|r| r.as_ref().err()))
25    {
26        Some(err) => Err(err.clone()),
27        None => Ok(()),
28    }
29}
30
31/// Unevaluated nodes reachable from `arrays`, inputs before consumers. Iterative, so graph
32/// depth is limited only by memory.
33fn topological_order(arrays: &[Array]) -> Vec<Arc<Node>> {
34    let mut visited = HashSet::new();
35    let mut order = Vec::new();
36    let mut stack: Vec<(Arc<Node>, bool)> =
37        arrays.iter().map(|a| (a.node.clone(), false)).collect();
38    while let Some((node, expanded)) = stack.pop() {
39        if expanded {
40            order.push(node);
41            continue;
42        }
43        if !visited.insert(Arc::as_ptr(&node)) {
44            continue;
45        }
46        let inputs: Vec<Arc<Node>> = match &*node.lock() {
47            State::Pending { inputs, .. } => inputs.iter().map(|a| a.node.clone()).collect(),
48            State::Done => continue,
49        };
50        stack.push((node, true));
51        stack.extend(
52            inputs
53                .into_iter()
54                .filter(|n| !visited.contains(&Arc::as_ptr(n)))
55                .map(|n| (n, false)),
56        );
57    }
58    order
59}
60
61/// Runs one node whose inputs are all evaluated, stores its outputs, and detaches it.
62fn evaluate(node: &Node) {
63    let mut state = node.lock();
64    let State::Pending { primitive, inputs } = std::mem::replace(&mut *state, State::Done) else {
65        return; // evaluated meanwhile by another thread
66    };
67    #[cfg(test)]
68    counter::inc();
69    match run(node, &primitive, &inputs) {
70        Ok(outputs) => {
71            for (output, data) in node.outputs.iter().zip(outputs) {
72                let _ = output.slot.set(Ok(data));
73            }
74        }
75        Err(err) => {
76            for output in &node.outputs {
77                let _ = output.slot.set(Err(err.clone()));
78            }
79        }
80    }
81    drop(state);
82    drop(inputs); // detach: release the inputs outside the lock
83}
84
85fn run(node: &Node, primitive: &Primitive, inputs: &[Array]) -> Result<Vec<Data>> {
86    let mut operands: Vec<Operand<'_>> = Vec::with_capacity(inputs.len());
87    for input in inputs {
88        match input
89            .out()
90            .slot
91            .get()
92            .expect("inputs are evaluated before consumers")
93        {
94            Ok(data) => operands.push((input.out(), data)),
95            Err(err) => return Err(err.clone()),
96        }
97    }
98    if node.stream.device.kind != DeviceKind::Cpu {
99        return Err(Error::UnsupportedDevice(node.stream.device));
100    }
101    let new = |buffer| Ok(vec![Data::new(buffer)]);
102    let share = |offset| {
103        Ok(vec![Data {
104            buffer: operands[0].1.buffer.clone(),
105            offset,
106        }])
107    };
108    match primitive {
109        Primitive::Unary(op) => new(kernels::unary(*op, operands[0])?),
110        Primitive::Binary(op) => new(kernels::binary(*op, operands[0], operands[1])?),
111        Primitive::Select => new(kernels::select(operands[0], operands[1], operands[2])?),
112        Primitive::AsType(dtype) => new(kernels::cast(operands[0], *dtype)?),
113        Primitive::Reduce { op, axes } => new(kernels::reduce(*op, operands[0], axes)?),
114        Primitive::ArgReduce { op, axis } => new(kernels::arg_reduce(*op, operands[0], *axis)?),
115        Primitive::Matmul => new(kernels::matmul(operands[0], operands[1])?),
116        Primitive::Transpose(_) | Primitive::Broadcast(_) => share(operands[0].1.offset),
117        Primitive::Slice { .. } => {
118            let (meta, data) = operands[0];
119            let (_, delta) =
120                view::slice_geometry(&meta.strides, &primitive.dim_slices(meta.shape.len())?);
121            share(data.offset + delta)
122        }
123        Primitive::Reshape(shape) => {
124            let (meta, data) = operands[0];
125            if view::nocopy_reshape(&meta.shape, &meta.strides, shape).is_some() {
126                share(data.offset)
127            } else {
128                new(kernels::contiguous(meta, data)?)
129            }
130        }
131        #[cfg(test)]
132        Primitive::AddMultiply => {
133            use crate::primitive::BinaryOp;
134            Ok(vec![
135                Data::new(kernels::binary(BinaryOp::Add, operands[0], operands[1])?),
136                Data::new(kernels::binary(
137                    BinaryOp::Multiply,
138                    operands[0],
139                    operands[1],
140                )?),
141            ])
142        }
143    }
144}
145
146/// Per-thread count of evaluated nodes, for run-once tests.
147#[cfg(test)]
148pub(crate) mod counter {
149    use std::cell::Cell;
150
151    thread_local!(static EVALUATED: Cell<usize> = const { Cell::new(0) });
152
153    pub(crate) fn inc() {
154        EVALUATED.set(EVALUATED.get() + 1);
155    }
156    pub(crate) fn take() -> usize {
157        EVALUATED.replace(0)
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use crate::buffer::live;
165
166    #[test]
167    fn shared_node_runs_once() {
168        let x = Array::from_slice(&[1.0f32, 2.0], &[2]).unwrap();
169        let shared = &x + &x;
170        let left = &shared * &x;
171        let right = &shared + &x;
172        let top = &left + &right;
173        counter::take();
174        eval(&[top.clone(), left.clone()]).unwrap();
175        assert_eq!(counter::take(), 4, "shared, left, right, top: once each");
176        top.eval().unwrap();
177        assert_eq!(counter::take(), 0, "already evaluated");
178        assert_eq!(top.to_vec::<f32>().unwrap(), vec![5.0, 14.0]);
179    }
180
181    #[test]
182    fn intermediates_freed() {
183        let before = live::count();
184        let mut x = Array::scalar(0i64);
185        let one = Array::scalar(1i64);
186        let steps = if cfg!(miri) { 100 } else { 10_000 };
187        for _ in 0..steps {
188            x = &x + &one;
189        }
190        x.eval().unwrap();
191        drop(one);
192        assert_eq!(
193            live::count() - before,
194            1,
195            "only the result's buffer is alive"
196        );
197        assert_eq!(x.item::<i64>().unwrap(), steps as i64);
198    }
199
200    #[test]
201    fn multi_output_once() {
202        let a = Array::from_slice(&[2i32, 3], &[2]).unwrap();
203        let b = Array::from_slice(&[5i32, 7], &[2]).unwrap();
204        let mut outs = Array::from_op(Primitive::AddMultiply, vec![a, b]).unwrap();
205        let product = outs.pop().unwrap();
206        let sum = outs.pop().unwrap();
207        counter::take();
208        eval(&[sum.clone(), product.clone()]).unwrap();
209        assert_eq!(counter::take(), 1);
210        drop(sum);
211        assert_eq!(product.to_vec::<i32>().unwrap(), vec![10, 21]);
212    }
213
214    #[test]
215    fn evaluation_detaches_inputs() {
216        let a = Array::scalar(1u8);
217        let c = &a + &a;
218        assert!(c.primitive().is_some());
219        c.eval().unwrap();
220        assert_eq!(c.primitive(), None);
221        assert_eq!(Arc::strong_count(&a.node), 1, "c no longer references a");
222    }
223}
nuvec_core/eval.rs

nuvec_core/
eval.rs
