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, Device, Error, Result, Stream, default_stream};
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        let leaves = operands
174            .iter()
175            .map(|(meta, data)| {
176                let buffer = crate::kernels::contiguous(meta, data)?;
177                Array::from_buffer(buffer, meta.dtype, &meta.shape)
178            })
179            .collect::<Result<Vec<_>>>()?;
180        let previous = default_stream();
181        crate::set_default_stream(Stream::default_for(Device::cpu()));
182        let outputs = self.expand(&leaves);
183        crate::set_default_stream(previous);
184        let outputs = outputs?;
185        crate::eval(&outputs)?;
186        outputs
187            .iter()
188            .map(|a| {
189                let data = a.data().expect("evaluated");
190                Ok(Data::new(crate::kernels::contiguous(a.out(), data)?))
191            })
192            .collect()
193    }
194
195    /// Reverse mode through the expanded operations, for the wanted float inputs.
196    pub(crate) fn vjp(
197        &self,
198        inputs: &[Array],
199        cotangents: &[Array],
200        wanted: &[bool],
201    ) -> Result<Vec<Option<Array>>> {
202        let active: Vec<usize> = (0..inputs.len())
203            .filter(|&i| wanted[i] && inputs[i].dtype().category() == Category::Float)
204            .collect();
205        let primals: Vec<Array> = active.iter().map(|&i| inputs[i].clone()).collect();
206        let f = |ps: &[Array]| self.expand(&substitute(inputs, &active, ps));
207        let (_, grads) = crate::vjp(f, &primals, cotangents)?;
208        let mut out = vec![None; inputs.len()];
209        for (g, &i) in grads.into_iter().zip(&active) {
210            out[i] = Some(g);
211        }
212        Ok(out)
213    }
214
215    /// Forward mode through the expanded operations (non-float inputs are constants).
216    pub(crate) fn jvp(&self, inputs: &[Array], tangents: &[Option<Array>]) -> Result<Vec<Array>> {
217        let active: Vec<usize> = (0..inputs.len())
218            .filter(|&i| tangents[i].is_some() && inputs[i].dtype().category() == Category::Float)
219            .collect();
220        let primals: Vec<Array> = active.iter().map(|&i| inputs[i].clone()).collect();
221        let ts: Vec<Array> = active
222            .iter()
223            .map(|&i| tangents[i].clone().expect("active"))
224            .collect();
225        if active.is_empty() {
226            return self
227                .expand(inputs)?
228                .iter()
229                .map(crate::derivatives::zeros_like)
230                .collect();
231        }
232        let f = |ps: &[Array]| self.expand(&substitute(inputs, &active, ps));
233        Ok(crate::jvp(f, &primals, &ts)?.1)
234    }
235
236    /// Batching rule: the expanded operations, batched.
237    pub(crate) fn vmap(
238        &self,
239        inputs: &[Array],
240        axes: &[Option<usize>],
241    ) -> Result<(Vec<Array>, usize)> {
242        let f = |xs: &[Array]| self.expand(xs);
243        Ok((crate::transforms::batch(&f, inputs, axes)?, 0))
244    }
245
246    /// The kernel name for inputs read as `loads`, with `per_thread` elements per thread.
247    pub(crate) fn kernel_name(&self, loads: &[Load], per_thread: usize) -> String {
248        let mut name = format!("nuvec_fused_{:016x}_", self.0.id);
249        name.extend(loads.iter().map(|l| match l {
250            Load::Scalar => 's',
251            Load::Flat => 'f',
252            Load::Strided(_) => 'g',
253        }));
254        let _ = write!(name, "_n{per_thread}");
255        name
256    }
257
258    /// Metal source of the kernel `name`, reading inputs as `loads`, `per_thread`
259    /// consecutive elements per thread. Arguments: the inputs, the outputs, the inputs'
260    /// element offsets, the (collapsed) shape, the strided inputs' strides (`ndim` each,
261    /// in order), `ndim` and the element count.
262    pub(crate) fn metal_source(&self, name: &str, loads: &[Load], per_thread: usize) -> String {
263        let g = &*self.0;
264        let (n, m) = (g.inputs.len(), g.outputs.len());
265        let mut s = String::from(
266            "#include \"common.h\"\nusing namespace metal;\nusing namespace nuvec;\n\n",
267        );
268        let _ = writeln!(s, "[[kernel]] void {name}(");
269        for (i, d) in g.inputs.iter().enumerate() {
270            let _ = writeln!(s, "    device const {}* in{i} [[buffer({i})]],", msl(*d));
271        }
272        for (k, d) in self.dtypes().into_iter().enumerate() {
273            let _ = writeln!(s, "    device {}* out{k} [[buffer({})]],", msl(d), n + k);
274        }
275        let _ = writeln!(s, "    constant long* offs [[buffer({})]],", n + m);
276        let _ = writeln!(s, "    constant int* shape [[buffer({})]],", n + m + 1);
277        let _ = writeln!(s, "    constant long* strides [[buffer({})]],", n + m + 2);
278        let _ = writeln!(s, "    constant int& ndim [[buffer({})]],", n + m + 3);
279        let _ = writeln!(s, "    constant uint& size [[buffer({})]],", n + m + 4);
280        s.push_str("    uint tid [[thread_position_in_grid]]) {\n");
281        for (i, d) in g.inputs.iter().enumerate() {
282            let _ = writeln!(s, "  device const {}* a{i} = in{i} + offs[{i}];", msl(*d));
283        }
284        let _ = writeln!(
285            s,
286            "  for (uint i = tid * {per_thread}; i < min(size, (tid + 1) * {per_thread}); ++i) {{"
287        );
288        for (i, (d, load)) in g.inputs.iter().zip(loads).enumerate() {
289            let index = match load {
290                Load::Scalar => "0".into(),
291                Load::Flat => "i".into(),
292                Load::Strided(k) => format!("loc(i, shape, strides + {k} * ndim, ndim)"),
293            };
294            let _ = writeln!(s, "    {} x{i} = a{i}[{index}];", msl(*d));
295        }
296        let arg = |r: &FusedRef| match *r {
297            FusedRef::Input(i) => format!("x{i}"),
298            FusedRef::Step(j) => format!("t{j}"),
299        };
300        let dtype = |r: &FusedRef| match *r {
301            FusedRef::Input(i) => g.inputs[i],
302            FusedRef::Step(j) => g.dtypes[j],
303        };
304        for (k, (primitive, args)) in g.steps.iter().enumerate() {
305            let a: Vec<String> = args.iter().map(arg).collect();
306            let expr = match primitive {
307                Primitive::Unary(op) => format!("op_{}({})", unary_fn(*op), a[0]),
308                Primitive::Binary(op) => binary_expr(*op, &a[0], &a[1]),
309                Primitive::Select => format!("({}.v ? {} : {})", a[0], a[1], a[2]),
310                Primitive::AsType(to) => {
311                    format!("convert<{}, {}>({})", msl(dtype(&args[0])), msl(*to), a[0])
312                }
313                _ => unreachable!("validated by Fused::new"),
314            };
315            let _ = writeln!(s, "    {} t{k} = {expr};", msl(g.dtypes[k]));
316        }
317        for (k, o) in g.outputs.iter().enumerate() {
318            let _ = writeln!(s, "    out{k}[i] = t{o};");
319        }
320        s.push_str("  }\n}\n");
321        s
322    }
323}
324
325/// `inputs` with the positions in `active` replaced by `values`.
326fn substitute(inputs: &[Array], active: &[usize], values: &[Array]) -> Vec<Array> {
327    let mut all = inputs.to_vec();
328    for (v, &i) in values.iter().zip(active) {
329        all[i] = v.clone();
330    }
331    all
332}
333
334/// The Metal type of a GPU dtype (see `common.h`).
335fn msl(d: DType) -> &'static str {
336    match d {
337        DType::Bool => "nbool",
338        DType::U8 => "uchar",
339        DType::U16 => "ushort",
340        DType::U32 => "uint",
341        DType::U64 => "ulong",
342        DType::I8 => "char",
343        DType::I16 => "short",
344        DType::I32 => "int",
345        DType::I64 => "long",
346        DType::F16 => "half",
347        DType::BF16 => "bfloat",
348        DType::F32 => "float",
349        DType::F64 | DType::Complex64 => unreachable!("not fused"),
350    }
351}
352
353fn unary_fn(op: UnaryOp) -> &'static str {
354    match op {
355        UnaryOp::Negative => "neg",
356        UnaryOp::Abs => "abs",
357        UnaryOp::Exp => "exp",
358        UnaryOp::Log => "log",
359        UnaryOp::Sqrt => "sqrt",
360        UnaryOp::Sin => "sin",
361        UnaryOp::Cos => "cos",
362        UnaryOp::Tanh => "tanh",
363        UnaryOp::Sigmoid => "sigmoid",
364        UnaryOp::ErfInv => "erfinv",
365        UnaryOp::Conjugate => unreachable!("not fused"),
366    }
367}
368
369/// The element-wise kernels' expressions (`elementwise.metal`).
370fn binary_expr(op: BinaryOp, a: &str, b: &str) -> String {
371    match op {
372        BinaryOp::Add => format!("op_add({a}, {b})"),
373        BinaryOp::Subtract => format!("op_sub({a}, {b})"),
374        BinaryOp::Multiply => format!("op_mul({a}, {b})"),
375        BinaryOp::Divide => format!("op_div({a}, {b})"),
376        BinaryOp::Maximum => format!("op_max({a}, {b})"),
377        BinaryOp::Minimum => format!("op_min({a}, {b})"),
378        BinaryOp::LogAddExp => format!("op_logaddexp({a}, {b})"),
379        BinaryOp::Equal => format!("nbool{{uchar(op_eq({a}, {b}))}}"),
380        BinaryOp::NotEqual => format!("nbool{{uchar(!op_eq({a}, {b}))}}"),
381        BinaryOp::Less => format!("nbool{{uchar(op_lt({a}, {b}))}}"),
382        BinaryOp::LessEqual => format!("nbool{{uchar(op_lt({a}, {b}) || op_eq({a}, {b}))}}"),
383        BinaryOp::Greater => format!("nbool{{uchar(op_lt({b}, {a}))}}"),
384        BinaryOp::GreaterEqual => format!("nbool{{uchar(op_lt({b}, {a}) || op_eq({a}, {b}))}}"),
385    }
386}
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391    use crate::ops;
392
393    fn sigmoid_mul() -> Fused {
394        // x * sigmoid(x) + y
395        Fused::new(
396            &[DType::F32, DType::F32],
397            vec![
398                (Primitive::Unary(UnaryOp::Sigmoid), vec![FusedRef::Input(0)]),
399                (
400                    Primitive::Binary(BinaryOp::Multiply),
401                    vec![FusedRef::Input(0), FusedRef::Step(0)],
402                ),
403                (
404                    Primitive::Binary(BinaryOp::Add),
405                    vec![FusedRef::Step(1), FusedRef::Input(1)],
406                ),
407            ],
408            &[1, 2],
409        )
410        .unwrap()
411    }
412
413    #[test]
414    fn validates_steps() {
415        let bad = |inputs: &[DType], steps| Fused::new(inputs, steps, &[0]).is_err();
416        assert!(bad(&[DType::F32], vec![]));
417        assert!(bad(
418            &[DType::F32],
419            vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Step(0)])]
420        ));
421        assert!(bad(
422            &[DType::I32],
423            vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Input(0)])]
424        ));
425        assert!(bad(
426            &[DType::F64],
427            vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Input(0)])]
428        ));
429        assert!(bad(&[DType::F32], vec![(Primitive::Matmul, vec![])]));
430        let f = sigmoid_mul();
431        assert_eq!((f.len(), f.dtypes()), (3, vec![DType::F32; 2]));
432        let one = vec![(Primitive::Unary(UnaryOp::Exp), vec![FusedRef::Input(0)])];
433        assert!(
434            Fused::new(&[DType::F32], one, &[1]).is_err(),
435            "no such output step"
436        );
437        assert_eq!(f, sigmoid_mul(), "equal graphs are equal");
438    }
439
440    fn values(arrays: &[Array]) -> Vec<Vec<f32>> {
441        arrays.iter().map(|a| a.to_vec::<f32>().unwrap()).collect()
442    }
443
444    /// Equal up to rounding (Miri perturbs transcendental functions by a few ulps).
445    fn assert_close(a: &[Vec<f32>], b: &[Vec<f32>]) {
446        assert_eq!(a.len(), b.len());
447        for (x, y) in a.iter().flatten().zip(b.iter().flatten()) {
448            assert!((x - y).abs() <= 1e-5 * x.abs().max(1.0), "{a:?} vs {b:?}");
449        }
450    }
451
452    #[test]
453    fn evaluates_like_its_operations_on_the_cpu() {
454        crate::set_default_stream(Stream::default_for(Device::cpu()));
455        let x = Array::from_slice(&[-2.0f32, 0.5, 3.0, 1.0, -1.0, 4.0], &[2, 3]).unwrap();
456        let y = Array::from_slice(&[1.0f32, 2.0], &[2, 1]).unwrap();
457        let y = ops::broadcast_to(&y, &[2, 3]).unwrap();
458        let xt = x.t().unwrap().t().unwrap();
459        let fused = sigmoid_mul();
460        let out = Array::from_op(Primitive::Compiled(fused.clone()), vec![xt, y.clone()]).unwrap();
461        assert_eq!(out.len(), 2);
462        assert_close(&values(&out), &values(&fused.expand(&[x, y]).unwrap()));
463    }
464
465    #[test]
466    fn differentiates_through_its_operations() {
467        crate::set_default_stream(Stream::default_for(Device::cpu()));
468        let fused = sigmoid_mul();
469        let x = Array::from_slice(&[-1.0f32, 0.0, 2.0], &[3]).unwrap();
470        let y = Array::from_slice(&[3.0f32, 1.0, 0.0], &[3]).unwrap();
471        let cot = Array::from_slice(&[1.0f32, 2.0, 0.5], &[3]).unwrap();
472        let cot2 = Array::from_slice(&[-1.0f32, 0.25, 3.0], &[3]).unwrap();
473        let run = |expanded: bool| {
474            let f = |xs: &[Array]| match expanded {
475                true => fused.expand(xs),
476                false => Array::from_op(Primitive::Compiled(fused.clone()), xs.to_vec()),
477            };
478            let primals = [x.clone(), y.clone()];
479            let (_, g) = crate::vjp(f, &primals, &[cot.clone(), cot2.clone()]).unwrap();
480            let (_, t) = crate::jvp(f, &primals, &[cot.clone(), cot2.clone()]).unwrap();
481            values(&[g, t].concat())
482        };
483        assert_close(&run(false), &run(true));
484        let batched = crate::vmap(
485            |xs| Array::from_op(Primitive::Compiled(fused.clone()), xs.to_vec()),
486            &[Some(0), None],
487            &[],
488        );
489        let xs = ops::reshape(&x, &[3, 1]).unwrap();
490        let ys = Array::from_slice(&[3.0f32], &[1]).unwrap();
491        let out = batched(&[xs.clone(), ys.clone()]).unwrap();
492        let expected = fused
493            .expand(&[xs, ops::broadcast_to(&ys, &[3, 1]).unwrap()])
494            .unwrap();
495        assert_close(&values(&out), &values(&expected));
496    }
497}
nuvec_core/fused.rs

nuvec_core/
fused.rs
