nuvec_core/
fused.rs

1//! Fused element-wise computations: one kernel for a graph of element-wise operations
2//! (MLX's `Compiled` primitive), as built by `compile`.
3//!
4//! On the GPU the graph becomes one generated Metal kernel that reads each input once and
5//! writes only the result. Everywhere else (and for the transforms) it expands into the
6//! operations it fuses.
7
8use std::fmt::Write as _;
9use std::hash::{DefaultHasher, Hash, Hasher};
10use std::sync::Arc;
11
12use crate::array::Data;
13use crate::kernels::Operand;
14use crate::primitive::{BinaryOp, Primitive, UnaryOp};
15use crate::{Array, Category, DType, Error, Result};
16
17/// An operand of a fused step: an input of the fused computation or an earlier step's
18/// result.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
20pub enum FusedRef {
21    /// The input at this position.
22    Input(usize),
23    /// The result of the step at this position.
24    Step(usize),
25}
26
27/// A graph of element-wise operations evaluated as one kernel, with one or more of its
28/// steps as results. Cheap to clone; equal graphs compare and hash equal.
29#[derive(Debug, Clone, PartialEq, Eq, Hash)]
30pub struct Fused(Arc<Graph>);
31
32#[derive(Debug, PartialEq, Eq, Hash)]
33struct Graph {
34    inputs: Vec<DType>,
35    steps: Vec<(Primitive, Vec<FusedRef>)>,
36    /// The steps whose results are outputs.
37    outputs: Vec<usize>,
38    /// The element type of each step's result.
39    dtypes: Vec<DType>,
40    /// A hash of the above, naming the generated kernels.
41    id: u64,
42}
43
44/// How a generated kernel reads one input.
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub(crate) enum Load {
47    /// One element, broadcast.
48    Scalar,
49    /// Laid out like the output.
50    Flat,
51    /// Through shape and strides (the `k`-th strided input).
52    Strided(usize),
53}
54
55fn computes(d: DType) -> bool {
56    !matches!(d, DType::F64 | DType::Complex64)
57}
58
59fn is_float(d: DType) -> bool {
60    matches!(d, DType::F16 | DType::BF16 | DType::F32)
61}
62
63impl Fused {
64    /// Whether `primitive` applied to inputs of these dtypes can be part of a fused
65    /// kernel: unary, binary, select and type conversions over the GPU's element types
66    /// (all but float64 and complex64).
67    pub fn supports(primitive: &Primitive, dtypes: &[DType]) -> bool {
68        if !dtypes.iter().copied().all(computes) {
69            return false;
70        }
71        match (primitive, dtypes) {
72            (Primitive::Unary(op), &[d]) => match op {
73                UnaryOp::Negative => d != DType::Bool,
74                UnaryOp::Abs => true,
75                UnaryOp::Conjugate => false,
76                _ => is_float(d),
77            },
78            (Primitive::Binary(op), &[a, b]) => {
79                a == b && (!matches!(op, BinaryOp::Divide | BinaryOp::LogAddExp) || is_float(a))
80            }
81            (Primitive::Select, &[c, a, b]) => c == DType::Bool && a == b,
82            (Primitive::AsType(to), &[_]) => computes(*to),
83            _ => false,
84        }
85    }
86
87    /// The fused computation of `steps` over inputs of the given dtypes, with the results
88    /// of the steps at positions `outputs` as its outputs. Each step is a primitive
89    /// [`Fused::supports`] and its operands, which refer to inputs or earlier steps.
90    pub fn new(
91        inputs: &[DType],
92        steps: Vec<(Primitive, Vec<FusedRef>)>,
93        outputs: &[usize],
94    ) -> Result<Self> {
95        let invalid = |what: String| Err(Error::InvalidArgument(format!("fused: {what}")));
96        if steps.is_empty() || outputs.is_empty() {
97            return invalid("no steps or no outputs".into());
98        }
99        if let Some(o) = outputs.iter().find(|&&o| o >= steps.len()) {
100            return invalid(format!("output step {o} of {}", steps.len()));
101        }
102        let mut dtypes: Vec<DType> = Vec::with_capacity(steps.len());
103        for (k, (primitive, args)) in steps.iter().enumerate() {
104            let mut ds = Vec::with_capacity(args.len());
105            for arg in args {
106                ds.push(match *arg {
107                    FusedRef::Input(i) if i < inputs.len() => inputs[i],
108                    FusedRef::Step(j) if j < k => dtypes[j],
109                    _ => return invalid(format!("step {k} refers to {arg:?}")),
110                });
111            }
112            if !Self::supports(primitive, &ds) {
113                return invalid(format!("{} on {ds:?} cannot be fused", primitive.name()));
114            }
115            dtypes.push(match primitive {
116                Primitive::Binary(op) if op.is_comparison() => DType::Bool,
117                Primitive::Select => ds[1],
118                Primitive::AsType(to) => *to,
119                _ => ds[0],
120            });
121        }
122        let mut graph = Graph {
123            inputs: inputs.to_vec(),
124            steps,
125            outputs: outputs.to_vec(),
126            dtypes,
127            id: 0,
128        };
129        let mut h = DefaultHasher::new();
130        (&graph.inputs, &graph.steps, &graph.outputs).hash(&mut h);
131        graph.id = h.finish();
132        Ok(Self(Arc::new(graph)))
133    }
134
135    /// The input dtypes.
136    pub fn inputs(&self) -> &[DType] {
137        &self.0.inputs
138    }
139
140    /// The number of fused operations.
141    pub fn len(&self) -> usize {
142        self.0.steps.len()
143    }
144
145    /// Always false: a fused computation has at least one step.
146    pub fn is_empty(&self) -> bool {
147        false
148    }
149
150    /// The outputs' dtypes.
151    pub fn dtypes(&self) -> Vec<DType> {
152        self.0.outputs.iter().map(|&o| self.0.dtypes[o]).collect()
153    }
154
155    /// The same computation recorded as ordinary operations.
156    pub(crate) fn expand(&self, inputs: &[Array]) -> Result<Vec<Array>> {
157        let mut results: Vec<Array> = Vec::with_capacity(self.0.steps.len());
158        for (primitive, args) in &self.0.steps {
159            let args = args
160                .iter()
161                .map(|r| match *r {
162                    FusedRef::Input(i) => inputs[i].clone(),
163                    FusedRef::Step(j) => results[j].clone(),
164                })
165                .collect();
166            results.push(Array::from_op(primitive.clone(), args)?.swap_remove(0));
167        }
168        Ok(self.0.outputs.iter().map(|&o| results[o].clone()).collect())
169    }
170
171    /// Evaluates the expanded operations on the CPU, for work the GPU does not take.
172    pub(crate) fn eval_cpu(&self, operands: &[Operand<'_>]) -> Result<Vec<Data>> {
173        crate::fast::imp::run_on_cpu(operands, |leaves| self.expand(leaves))
174    }
175
176    /// Reverse mode through the expanded operations, for the wanted float inputs.
177    pub(crate) fn vjp(
178        &self,
179        inputs: &[Array],
180        cotangents: &[Array],
181        wanted: &[bool],
182    ) -> Result<Vec<Option<Array>>> {
183        let active: Vec<usize> = (0..inputs.len())
184            .filter(|&i| wanted[i] && inputs[i].dtype().category() == Category::Float)
185            .collect();
186        let primals: Vec<Array> = active.iter().map(|&i| inputs[i].clone()).collect();
187        let f = |ps: &[Array]| self.expand(&substitute(inputs, &active, ps));
188        let (_, grads) = crate::vjp_flat(f, &primals, cotangents)?;
189        let mut out = vec![None; inputs.len()];
190        for (g, &i) in grads.into_iter().zip(&active) {
191            out[i] = Some(g);
192        }
193        Ok(out)
194    }
195
196    /// Forward mode through the expanded operations (non-float inputs are constants).
197    pub(crate) fn jvp(&self, inputs: &[Array], tangents: &[Option<Array>]) -> Result<Vec<Array>> {
198        let active: Vec<usize> = (0..inputs.len())
199            .filter(|&i| tangents[i].is_some() && inputs[i].dtype().category() == Category::Float)
200            .collect();
201        let primals: Vec<Array> = active.iter().map(|&i| inputs[i].clone()).collect();
202        let ts: Vec<Array> = active
203            .iter()
204            .map(|&i| tangents[i].clone().expect("active"))
205            .collect();
206        if active.is_empty() {
207            return self
208                .expand(inputs)?
209                .iter()
210                .map(crate::derivatives::zeros_like)
211                .collect();
212        }
213        let f = |ps: &[Array]| self.expand(&substitute(inputs, &active, ps));
214        Ok(crate::jvp_flat(f, &primals, &ts)?.1)
215    }
216
217    /// Batching rule: the expanded operations, batched.
218    pub(crate) fn vmap(
219        &self,
220        inputs: &[Array],
221        axes: &[Option<usize>],
222    ) -> Result<(Vec<Array>, usize)> {
223        let f = |xs: &[Array]| self.expand(xs);
224        Ok((crate::transforms::batch(&f, inputs, axes)?, 0))
225    }
226
227    /// The kernel name for inputs read as `loads`, with `per_thread` elements per thread.
228    pub(crate) fn kernel_name(&self, loads: &[Load], per_thread: usize) -> String {
229        let mut name = format!("nuvec_fused_{:016x}_", self.0.id);
230        name.extend(loads.iter().map(|l| match l {
231            Load::Scalar => 's',
232            Load::Flat => 'f',
233            Load::Strided(_) => 'g',
234        }));
235        let _ = write!(name, "_n{per_thread}");
236        name
237    }
238
239    /// Metal source of the kernel `name`, reading inputs as `loads`, `per_thread`
240    /// consecutive elements per thread. Arguments: the inputs, the outputs, the inputs'
241    /// element offsets, the (collapsed) shape, the strided inputs' strides (`ndim` each,
242    /// in order), `ndim` and the element count.
243    pub(crate) fn metal_source(&self, name: &str, loads: &[Load], per_thread: usize) -> String {
244        let g = &*self.0;
245        let (n, m) = (g.inputs.len(), g.outputs.len());
246        let mut s = String::from(
247            "#include \"common.h\"\nusing namespace metal;\nusing namespace nuvec;\n\n",
248        );
249        let _ = writeln!(s, "[[kernel]] void {name}(");
250        for (i, d) in g.inputs.iter().enumerate() {
251            let _ = writeln!(s, "    device const {}* in{i} [[buffer({i})]],", msl(*d));
252        }
253        for (k, d) in self.dtypes().into_iter().enumerate() {
254            let _ = writeln!(s, "    device {}* out{k} [[buffer({})]],", msl(d), n + k);
255        }
256        let _ = writeln!(s, "    constant long* offs [[buffer({})]],", n + m);
257        let _ = writeln!(s, "    constant int* shape [[buffer({})]],", n + m + 1);
258        let _ = writeln!(s, "    constant long* strides [[buffer({})]],", n + m + 2);
259        let _ = writeln!(s, "    constant int& ndim [[buffer({})]],", n + m + 3);
260        let _ = writeln!(s, "    constant uint& size [[buffer({})]],", n + m + 4);
261        s.push_str("    uint tid [[thread_position_in_grid]]) {\n");
262        for (i, d) in g.inputs.iter().enumerate() {
263            let _ = writeln!(s, "  device const {}* a{i} = in{i} + offs[{i}];", msl(*d));
264        }
265        let _ = writeln!(
266            s,
267            "  for (uint i = tid * {per_thread}; i < min(size, (tid + 1) * {per_thread}); ++i) {{"
268        );
269        for (i, (d, load)) in g.inputs.iter().zip(loads).enumerate() {
270            let index = match load {
271                Load::Scalar => "0".into(),
272                Load::Flat => "i".into(),
273                Load::Strided(k) => format!("loc(i, shape, strides + {k} * ndim, ndim)"),
274            };
275            let _ = writeln!(s, "    {} x{i} = a{i}[{index}];", msl(*d));
276        }
277        let arg = |r: &FusedRef| match *r {
278            FusedRef::Input(i) => format!("x{i}"),
279            FusedRef::Step(j) => format!("t{j}"),
280        };
281        let dtype = |r: &FusedRef| match *r {
282            FusedRef::Input(i) => g.inputs[i],
283            FusedRef::Step(j) => g.dtypes[j],
284        };
285        for (k, (primitive, args)) in g.steps.iter().enumerate() {
286            let a: Vec<String> = args.iter().map(arg).collect();
287            let expr = match primitive {
288                Primitive::Unary(op) => format!("op_{}({})", unary_fn(*op), a[0]),
289                Primitive::Binary(op) => binary_expr(*op, &a[0], &a[1]),
290                Primitive::Select => format!("({}.v ? {} : {})", a[0], a[1], a[2]),
291                Primitive::AsType(to) => {
292                    format!("convert<{}, {}>({})", msl(dtype(&args[0])), msl(*to), a[0])
293                }
294                _ => unreachable!("validated by Fused::new"),
295            };
296            let _ = writeln!(s, "    {} t{k} = {expr};", msl(g.dtypes[k]));
297        }
298        for (k, o) in g.outputs.iter().enumerate() {
299            let _ = writeln!(s, "    out{k}[i] = t{o};");
300        }
301        s.push_str("  }\n}\n");
302        s
303    }
304}
305
306/// `inputs` with the positions in `active` replaced by `values`.
307fn substitute(inputs: &[Array], active: &[usize], values: &[Array]) -> Vec<Array> {
308    let mut all = inputs.to_vec();
309    for (v, &i) in values.iter().zip(active) {
310        all[i] = v.clone();
311    }
312    all
313}
314
315/// The Metal type of a GPU dtype (see `common.h`).
316fn msl(d: DType) -> &'static str {
317    match d {
318        DType::Bool => "nbool",
319        DType::U8 => "uchar",
320        DType::U16 => "ushort",
321        DType::U32 => "uint",
322        DType::U64 => "ulong",
323        DType::I8 => "char",
324        DType::I16 => "short",
325        DType::I32 => "int",
326        DType::I64 => "long",
327        DType::F16 => "half",
328        DType::BF16 => "bfloat",
329        DType::F32 => "float",
330        DType::F64 | DType::Complex64 => unreachable!("not fused"),
331    }
332}
333
334fn unary_fn(op: UnaryOp) -> &'static str {
335    match op {
336        UnaryOp::Negative => "neg",
337        UnaryOp::Abs => "abs",
338        UnaryOp::Exp => "exp",
339        UnaryOp::Log => "log",
340        UnaryOp::Sqrt => "sqrt",
341        UnaryOp::Sin => "sin",
342        UnaryOp::Cos => "cos",
343        UnaryOp::Tanh => "tanh",
344        UnaryOp::Sigmoid => "sigmoid",
345        UnaryOp::ErfInv => "erfinv",
346        UnaryOp::Conjugate => unreachable!("not fused"),
347    }
348}
349
350/// The element-wise kernels' expressions (`elementwise.metal`).
351fn binary_expr(op: BinaryOp, a: &str, b: &str) -> String {
352    match op {
353        BinaryOp::Add => format!("op_add({a}, {b})"),
354        BinaryOp::Subtract => format!("op_sub({a}, {b})"),
355        BinaryOp::Multiply => format!("op_mul({a}, {b})"),
356        BinaryOp::Divide => format!("op_div({a}, {b})"),
357        BinaryOp::Maximum => format!("op_max({a}, {b})"),
358        BinaryOp::Minimum => format!("op_min({a}, {b})"),
359        BinaryOp::LogAddExp => format!("op_logaddexp({a}, {b})"),
360        BinaryOp::Equal => format!("nbool{{uchar(op_eq({a}, {b}))}}"),
361        BinaryOp::NotEqual => format!("nbool{{uchar(!op_eq({a}, {b}))}}"),
362        BinaryOp::Less => format!("nbool{{uchar(op_lt({a}, {b}))}}"),
363        BinaryOp::LessEqual => format!("nbool{{uchar(op_lt({a}, {b}) || op_eq({a}, {b}))}}"),
364        BinaryOp::Greater => format!("nbool{{uchar(op_lt({b}, {a}))}}"),
365        BinaryOp::GreaterEqual => format!("nbool{{uchar(op_lt({b}, {a}) || op_eq({a}, {b}))}}"),
366    }
367}
368
369#[cfg(test)]
370mod tests {
371    use super::*;
372    use crate::ops::imp as ops;
373
374    fn sigmoid_mul() -> Fused {
375        // x * sigmoid(x) + y
376        Fused::new(
377            &[DType::F32, DType::F32],
378            vec![
379                (Primitive::Unary(UnaryOp::Sigmoid), vec![FusedRef::Input(0)]),
380                (
381                    Primitive::Binary(BinaryOp::Multiply),
382                    vec![FusedRef::Input(0), FusedRef::Step(0)],
383                ),
384                (
385                    Primitive::Binary(BinaryOp::Add),
386                    vec![FusedRef::Step(1), FusedRef::Input(1)],
387                ),
388            ],
389            &[1, 2],
390        )
391        .unwrap()
392    }
393
394    #[test]
395    fn validates_steps() {
396        let bad = |inputs: &[DType], steps| Fused::new(inputs, steps, &[0]).is_err();
397        assert!(bad(&[DType::F32], vec![]));
398        assert!(bad(
399            &[DType::F32],
400            vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Step(0)])]
401        ));
402        assert!(bad(
403            &[DType::I32],
404            vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Input(0)])]
405        ));
406        assert!(bad(
407            &[DType::F64],
408            vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Input(0)])]
409        ));
410        assert!(bad(&[DType::F32], vec![(Primitive::Matmul, vec![])]));
411        let f = sigmoid_mul();
412        assert_eq!((f.len(), f.dtypes()), (3, vec![DType::F32; 2]));
413        let one = vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Input(0)])];
414        assert!(
415            Fused::new(&[DType::F32], one, &[1]).is_err(),
416            "no such output step"
417        );
418        assert_eq!(f, sigmoid_mul(), "equal graphs are equal");
419    }
420
421    fn values(arrays: &[Array]) -> Vec<Vec<f32>> {
422        arrays.iter().map(|a| a.to_vec::<f32>().unwrap()).collect()
423    }
424
425    /// Equal up to rounding (Miri perturbs transcendental functions by a few ulps).
426    fn assert_close(a: &[Vec<f32>], b: &[Vec<f32>]) {
427        assert_eq!(a.len(), b.len());
428        for (x, y) in a.iter().flatten().zip(b.iter().flatten()) {
429            assert!((x - y).abs() <= 1e-5 * x.abs().max(1.0), "{a:?} vs {b:?}");
430        }
431    }
432
433    #[test]
434    fn evaluates_like_its_operations_on_the_cpu() {
435        crate::set_default_stream(crate::Stream::default_for(crate::Device::cpu()));
436        let x = Array::from_slice(&[-2.0f32, 0.5, 3.0, 1.0, -1.0, 4.0], [2, 3]);
437        let y = Array::from_slice(&[1.0f32, 2.0], [2, 1]);
438        let y = ops::broadcast_to(&y, &[2, 3]).unwrap();
439        let xt = x.t().t();
440        let fused = sigmoid_mul();
441        let out = Array::from_op(Primitive::Compiled(fused.clone()), vec![xt, y.clone()]).unwrap();
442        assert_eq!(out.len(), 2);
443        assert_close(&values(&out), &values(&fused.expand(&[x, y]).unwrap()));
444    }
445
446    #[test]
447    fn differentiates_through_its_operations() {
448        crate::set_default_stream(crate::Stream::default_for(crate::Device::cpu()));
449        let fused = sigmoid_mul();
450        let x = Array::from_slice(&[-1.0f32, 0.0, 2.0], [3]);
451        let y = Array::from_slice(&[3.0f32, 1.0, 0.0], [3]);
452        let cot = Array::from_slice(&[1.0f32, 2.0, 0.5], [3]);
453        let cot2 = Array::from_slice(&[-1.0f32, 0.25, 3.0], [3]);
454        let run = |expanded: bool| {
455            let f = |xs: &[Array]| match expanded {
456                true => fused.expand(xs),
457                false => Array::from_op(Primitive::Compiled(fused.clone()), xs.to_vec()),
458            };
459            let primals = [x.clone(), y.clone()];
460            let (_, g) = crate::vjp_flat(f, &primals, &[cot.clone(), cot2.clone()]).unwrap();
461            let (_, t) = crate::jvp_flat(f, &primals, &[cot.clone(), cot2.clone()]).unwrap();
462            values(&[g, t].concat())
463        };
464        assert_close(&run(false), &run(true));
465        let batched = crate::vmap_flat(
466            |xs| Array::from_op(Primitive::Compiled(fused.clone()), xs.to_vec()),
467            &[Some(0), None],
468            &[],
469        );
470        let xs = ops::reshape(&x, &[3, 1]).unwrap();
471        let ys = Array::from_slice(&[3.0f32], [1]);
472        let out = batched(&[xs.clone(), ys.clone()]).unwrap();
473        let expected = fused
474            .expand(&[xs, ops::broadcast_to(&ys, &[3, 1]).unwrap()])
475            .unwrap();
476        assert_close(&values(&out), &values(&expected));
477    }
478}
nuvec_core/fused.rs

nuvec_core/
fused.rs
